Sparse Polynomials: `boundlab.polysp`#

The idea#

A zonotope must give up exactness at every product: eps_i * eps_j is not affine, so it collapses into fresh noise. The sparse polynomial domain keeps those terms instead — a value is a sparse multivariable polynomial over error symbols:

\[ \textbf{x}_j \;=\; \sum_{-1 \le i_1, \cdots, i_d \le k} G_{j, i_1, \cdots, i_d} \varepsilon_{i_1} \cdots \varepsilon_{i_d} \text{ where } \varepsilon_{-1} = 1, \varepsilon_n \in [-1, 1] \text{ for } 0 \le n \le k \]

Here \(\textbf{x}_j\) is one element of the abstract value (the index \(j\) ranges over the tensor’s shape, so each element carries its own polynomial over the shared symbols \(\varepsilon_0, \ldots, \varepsilon_k\)). A multi-index \((i_1, \ldots, i_d)\) names one monomial \(\varepsilon_{i_1} \cdots \varepsilon_{i_d}\), where \(d\) — the order of the representation — caps the degree:

The sentinel \(\varepsilon_{-1} = 1\) makes lower-degree monomials a special case of degree-\(d\) ones — padding a slot with \(-1\) just drops a factor, so a monomial’s true degree is its count of non-sentinel slots. The all-\((-1)\) multi-index is the constant term, and a single non-\((-1)\) slot gives the affine (zonotope-like) part.
Repeated indices are meaningful: \((i, i)\) is \(\varepsilon_i^2\), which is how squares survive a product instead of collapsing into noise.
Since multiplication is commutative, only the sorted multi-index matters; the implementation keeps indices sorted so equal monomials can be recognized and their coefficients summed (coalescing).

The coefficient tensor \(G\) is sparse over the monomial axes: almost all multi-indices have zero coefficient, so a Generator stores only the nse nonzero monomials as parallel arrays — \(\mathrm{data}[r]\), a dense coefficient tensor of the value’s shape, and \(\mathrm{indices}[r]\), the -1-padded sorted multi-index of monomial \(r\) (row 0 is always the constant term). Concretization bounds every non-constant monomial by \([-1, 1]\):

\[ \textbf{x}_j \;\in\; \Big[\, \mathrm{data}_{0j} - \sum_{r \ge 1} |\mathrm{data}_{rj}|,\;\; \mathrm{data}_{0j} + \sum_{r \ge 1} |\mathrm{data}_{rj}| \,\Big], \]

sound always, and tight when the monomials are independent.

TODO: Improving this bound using Legendre polynomial.

Handling Basic Operators#

Linear operators never touch the monomial structure: add, sub, neg, reshape, transpose, reductions, and any product with a constant (Gemm, x @ W) act on the coefficient tensor \(G\) alone, one exact einsum carried across the monomial rows. Adding two values with different symbol tables first remaps their indices into the aligned table, then coalesces rows with identical monomials (sparse_poly.sum).

Non-linear activations (relu, exp, tanh, reciprocal) reuse the zonotope linearizers: f(x) ⊆ slope·x + bias ± error, where the affine part is exact on the polynomial and the residual arrives as fresh Noise. The after_each normalizer keep_ploysp immediately lifts that noise into a new order-1 symbol, so later products and sums can still cancel against it.

Handling `Matmul`#

A product of two sparse polynomials multiplies every monomial pair — \((d_r m_r)(e_s m_s) = (d_r e_s)\, m_r m_s\) — an nse₁ × nse₂ grid of small matrix products. sparse_poly.matmul keeps that grid tractable in two truncation steps, each of which encloses everything it drops in interval noise (sound because every monomial ranges within \([-1, 1]\)):

pre_matmul — drop input monomials whose coefficient (Frobenius) norm falls below a threshold derived from eps, before the grid is formed. Writing gen = kept + dropped, the discarded part of the product is enclosed by n₁@|g₂| + |g₁|@n₂ + n₁@n₂ with nᵢ the summed dropped magnitudes.
final_matmul — form the surviving pair products and keep only those whose magnitude clears eps and whose combined degree fits order; the rest are summed into noise. Kept pairs get the sorted concatenation of their index rows as their monomial, and duplicates are coalesced.

On CUDA a Triton kernel walks the pair grid in tiles, filtering and compacting on the fly, so nothing of size nse₁ · nse₂ is ever materialized; elsewhere the reference implementation applies the same filters. sparse_applypoly runs the same pipeline for the elementwise quadratic (a·x + b)·x.

Handling `Softmax`#

The default path is the shared decomposition \(\sigma_i = 1 / \sum_j e^{\nu_j - \nu_i}\) (Softmax2ExpReciprocal): pairwise differences and the reduce-sum are exact on polynomials, and exp / reciprocal use the linearizers above.

SoftmaxConstrained (boundlab.polysp.softmax) tightens the result with the simplex constraint. Softmax outputs satisfy \(\sum_i \sigma_i(\nu) = 1\) identically, so the propagated abstract row sums

\[ E_q(\varepsilon) \;=\; \sum_i \hat\sigma_{qi}(\varepsilon) - 1 \]

are polynomials known to evaluate to zero at every feasible \(\varepsilon\) — pure representation slack. Consequently, for any multipliers \(t\),

\[ \hat\sigma_{qi}' \;=\; \hat\sigma_{qi} + t_{qi}\, E_q \]

represents the same function, and softmax_opt chooses the \(t\) that minimizes the concretized width in three steps:

Shrink pinned symbols. Each equation with a nonzero linear coefficient \(a_{qj}\) on symbol \(\varepsilon_j\) pins it to

\[ \varepsilon_j \;\in\; -\frac{c_q}{a_{qj}} \pm \frac{\rho_q}{|a_{qj}|}, \]

where \(c_q\) is the equation’s constant and \(\rho_q\) encloses its other linear and (concretized) higher-order terms; intersecting over the equations and with \([-1, 1]\) gives a range \(c_j \pm h_j\) (bounding_error_terms). Substituting \(\varepsilon_j = s_j\, \varepsilon_j'\) with \(s_j = \min(|c_j| + h_j,\, 1)\) is sound and rescales every monomial coefficient accordingly (restrict_errors).
Least-squares direction. Let \(A\) and \(B\) collect the non-constant monomial rows of the equations and of the output (the constant row is skipped — only error rows contribute width). Solve \(T = \arg\min_T \lVert A\,T + B \rVert_2\) column-wise.
Exact L1 step. The concretized width is the L1 norm of the error rows, so rescale the direction by the per-column step

\[ s^\ast = \arg\min_s\; \lVert B + s\,(A\,T) \rVert_1 \;=\; \operatorname{median}_w\Big( -\tfrac{B_r}{(AT)_r} \Big), \quad w_r = |(A\,T)_r| , \]

the weighted median of the breakpoints of this piecewise-linear objective. \(s = 0\) is always a candidate, so the correction never increases the width.

The constraint’s error rows are anti-correlated with the output’s — that is what the least-squares fit finds — so adding the right multiple cancels error mass that no per-operator relaxation could remove.

Compression#

Products grow the monomial count multiplicatively: after one attention layer of the SST comparison the value holds ~360k monomials, and the pair grid of the next attention layer — not precision — becomes the binding constraint. Compression, proposed in Future Plans for polysp §2 and now implemented in this module, periodically re-expresses a value over fewer, learned symbols, paying a certified noise term for the modeling error.

Warning

Compression is not fully tested yet. Soundness is unit-tested by sampling on synthetic values, and the solver settings were tuned on real layer-0 barrier matrices of the SST comparison — but no complete end-to-end verified run has exercised it (the layer-2 attn @ V pair product still exhausts memory in the reference kernel before the run finishes), so its effect on final network bounds is unvalidated.

The key observation: every non-constant monomial \(\mu_r\) ranges in \([-1, 1]\), so a value is a linear map of its monomial vector, \(x = d_0 + A\mu\) with one coefficient column per monomial. PolySp.compress factorizes \(A \approx D S\) with nterms - 1 dictionary atoms by LAD-lasso dictionary learning (\(\min \|A - DS\|_1 + \lambda\|S\|_1\), atoms in the unit ball), solved by lasso.lad.admm_dict_learning from the pytorch-lasso package. Each row of \(S\) then defines a fresh error symbol

\[ \varepsilon_j' \;=\; \frac{S_{j\cdot}\,\mu}{\lVert S_{j\cdot}\rVert_1} \;\in\; [-1, 1], \]

carried with amplitude \(\lVert S_{j\cdot}\rVert_1 D_j\) — sound, because the true \(\varepsilon'\) set is a subset of the box — while the factorization error and the pre-filtered tail are enclosed coefficient-wise in a returned Noise (\(\sum_r |A - DS|_{ir}\) per element, exactly the LAD fit term the solver minimizes).

Two choices decide whether this pays off, both measured on real barrier matrices:

Residual mass must stay tiny. Generator rows propagate through a downstream linear layer as \(W d\) (signs cancel); Noise propagates as \(|W|\,h\) (they cannot). A cold-started solve leaves ~90% of the mass in the residual and inflates downstream width ~7.7x.
Warm start at the truncation solution. The solver seeds \(D\) with the heaviest monomial columns and \(S\) with the matching diagonal, so iteration 0 is top-nterms truncation and every step moves tail mass from noise into symbols. On the 361k-monomial attention barrier the defaults reach ~1.02x mean downstream width with ~0.6% noise in ~15 s (raising max_iter to 50 with the inner tolerance disabled reaches ~1.01x mean / ~1.2x worst-element at ~140 s).

Compression triggers at layer barriers: the model marks values with boundlab.ops.marked_idenity(x, name="layer_barrier") (a custom MarkedIdentity ONNX node), and the PolySpCompression handler compresses any marked value whose term count exceeds its nterm_limit, adding the certified residual back through keep_ploysp. The caveat from the plan remains: fresh symbols forget their dependence on the old ones, so compressing one value severs its correlation with other live values — compressing all live values of a layer jointly is future work.

`interpret`#

The sparse polynomial interpreter, exactly as assembled in boundlab.polysp:

from boundlab.ibp.softmax import Softmax2ExpReciprocal
from boundlab.interp import Interpreter, base
from boundlab.polysp import Matmul, PolySpCompression, keep_ploysp
from boundlab.zono import linearizers

interpret = Interpreter(
    base.interpret,                    # shared exact operators
    Matmul(eps=1e-4, order=4),         # sparse pair product (see above)
    linearizers.MaxWithConst2Relu(),
    linearizers.Relu(),                # zonotope linearizers; residuals are
    linearizers.Tanh(),                # lifted back into polynomial symbols
    linearizers.Exp(),
    linearizers.Reciprocal(),
    Softmax2ExpReciprocal(),
    PolySpCompression(),               # compress at layer barriers (see above)
    after_each=keep_ploysp,            # Bias/Noise -> PolySp
)

Swap in SoftmaxConstrained for the simplex-tightened softmax, or a Poly(eps=..., order=...) handler to apply concrete quadratics through the sparse pair product. Abstract inputs are built with PolySp.error(err) * radius + center; the PolySp/Generator API details are in the API reference.

Sparse Polynomials: boundlab.polysp#