Future Plans for polysp#

Note

This page is a research proposal: none of it is implemented yet. It formalizes three directions for the sparse polynomial domain and records the questions each one has to answer before it can land.

Throughout, a value is the sparse polynomial of Sparse Polynomials: boundlab.polysp,

\[ \textbf{x}_j(\varepsilon) \;=\; \sum_{-1 \le i_1, \cdots, i_d \le k} G_{j, i_1, \cdots, i_d}\, \varepsilon_{i_1} \cdots \varepsilon_{i_d}, \qquad \varepsilon \in \Omega = [-1, 1]^k, \]

and we equip \(\Omega\) with the uniform probability measure \(\mathrm{d}\mu = 2^{-k}\, \mathrm{d}\varepsilon\) and the inner product

\[ \langle f, g \rangle \;=\; \int_\Omega f(\varepsilon)\, g(\varepsilon)\, \mathrm{d}\mu(\varepsilon). \]

This inner product defines the Legendre basis of §1 and prices the approximation error of §3; §2 works directly on coefficient tensors.

1. Legendre Polynomial Concretization#

Motivation#

Concretization currently bounds every monomial independently by \([-1, 1]\): \(\textbf{x}_j \in \mathrm{data}_{0j} \pm \sum_{r \ge 1} |\mathrm{data}_{rj}|\). This ignores that all monomials are functions of the same \(\varepsilon\). Two elementary examples of the slack:

  • \(\varepsilon^2\) has true range \([0, 1]\); the monomial bound reports \([-1, 1]\) — twice the width, wrongly centered.

  • \(\varepsilon - \varepsilon^3\) has true range \([-\tfrac{2}{3\sqrt{3}}, \tfrac{2}{3\sqrt{3}}] \approx [-0.385, 0.385]\); the monomial bound reports \([-2, 2]\) — a factor of five, because the two monomials never attain their extremes at the same \(\varepsilon\).

Background: Legendre polynomials#

No prior familiarity assumed. Apply Gram–Schmidt to the monomials \(1, t, t^2, t^3, \ldots\) under the inner product \(\langle f, g\rangle = \tfrac12 \int_{-1}^{1} f g\, \mathrm{d}t\). The result is the Legendre family — the polynomials that are pairwise orthogonal on \([-1, 1]\):

\[ P_0 = 1, \quad P_1 = t, \quad P_2 = \tfrac{3t^2 - 1}{2}, \quad P_3 = \tfrac{5t^3 - 3t}{2}, \qquad \langle P_i, P_j \rangle = \frac{\delta_{ij}}{2i + 1}. \]

Three facts matter here:

  1. Same span. \(P_0, \ldots, P_m\) span exactly the polynomials of degree \(\le m\), so rewriting a polynomial in this basis is a lossless, triangular change of coordinates.

  2. Unit range. \(|P_i(t)| \le 1\) on \([-1, 1]\), with \(P_i(1) = 1\) — each basis function is as “monomial-like” as possible for bounding purposes.

  3. Mean extraction. Orthogonality against \(P_0 = 1\) means every \(P_i\) with \(i \ge 1\) has zero mean; the \(P_0\)-coefficient of a function is its average over the interval.

In \(k\) variables the products \(P_\alpha(\varepsilon) = \prod_{n} P_{\alpha_n}(\varepsilon_n)\) (multi-index \(\alpha\)) inherit all three properties with respect to \(\langle\cdot,\cdot\rangle\) on \(\Omega\).

Proposal#

Convert the monomial representation to the Legendre basis before concretizing. Each row’s monomial \(\prod \varepsilon_{i}\) factors into per-symbol powers \(\varepsilon_n^{m_n}\); expanding each power in \(P_0, \ldots, P_{m_n}\) and multiplying out rewrites the value as

\[ \textbf{x}_j \;=\; \sum_\alpha \ell_{j\alpha}\, P_\alpha(\varepsilon), \qquad\text{and then}\qquad \textbf{x}_j \;\in\; \ell_{j0} \pm \sum_{\alpha \ne 0} |\ell_{j\alpha}| \]

is sound by the unit-range property — the same shape of bound as before, but now centered at the value’s true mean \(\ell_{j0}\), with the basis change absorbing the dependencies between powers of a shared symbol. On the examples above: \(\varepsilon^2 = \tfrac13 P_0 + \tfrac23 P_2\) gives \([-\tfrac13, 1]\) (vs. \([-1, 1]\)), and \(\varepsilon - \varepsilon^3 = \tfrac25 P_1 - \tfrac25 P_3\) gives \([-0.8, 0.8]\) (vs. \([-2, 2]\)).

Considerations#

  • Sparsity growth. A monomial of degree \(m\) expands into at most \(\prod_n (m_n + 1)\) Legendre terms; with the order capped at \(d\) this is a bounded (\(\le 2^d\)-fold) blow-up, and terms coalesce. Still, whether to convert only inside lbub() or to maintain the Legendre form natively (products of Legendre polynomials re-expand by known linearization formulas) is an open cost trade-off.

  • Remaining slack. The bound is still not exact: all \(P_\alpha\) attain \(+1\) simultaneously at the all-ones corner, but the signs of \(\ell_\alpha\) may prevent the bound from being attained. Exactness holds elementwise for univariate values of degree \(\le 1\) only; quantify the gap empirically against sampling.

  • Attribution. Reasons breakdowns currently follow monomial rows; the basis change mixes rows of one symbol but never across symbols, so per-symbol attribution survives — per-reason weights need a mapping rule.

2. Polynomial Compression / Tensor Decomposition#

Motivation#

Products grow the monomial count multiplicatively; after a few attention layers the nse axis, not precision, is the binding constraint. The proposal: periodically re-express a value over fewer, learned symbols, paying a certified noise term for the modeling error.

Warm-up: the zonotope (degree-1) case#

A zonotope block is a matrix \(G \in \mathbb{R}^{p \times k}\). Compress it with a LASSO-style factorization

\[ \min_{G', S}\; \|G - G' S\|_2^2 \;+\; \lambda \sum_{j} \|S_{j\cdot}\|_1, \qquad G' \in \mathbb{R}^{p \times k'},\; S \in \mathbb{R}^{k' \times k},\; k' \ll k . \]

Why the \(L^1\) penalty is the right one. Define candidate new symbols \(\varepsilon'_j = S_{j\cdot}\, \varepsilon\). For \(\varepsilon \in [-1,1]^k\),

\[ |\varepsilon'_j| \;\le\; \|S_{j\cdot}\|_1 , \]

so whenever \(\|S_{j\cdot}\|_1 \le 1\) (achievable w.l.o.g. by rescaling the row into \(G'\)), treating \(\varepsilon'_j\) as a fresh symbol in \([-1, 1]\) is sound. Then

\[ G\varepsilon \;=\; G'\varepsilon' + (G - G'S)\,\varepsilon \;\subseteq\; G'\varepsilon' \;\pm\; \textstyle\sum_k |G - G'S|_{\cdot k}, \]

i.e. the compressed value is \(G'\) over \(k'\) symbols plus a residual noise term read off the \(L^1\) row norms of the factorization error. The penalty therefore does double duty: it sparsifies \(S\) and keeps the new symbols’ ranges — hence the soundness rescaling — from inflating \(G'\).

Solving the LASSO: ISTA#

The standard solver for objectives of the form \(\min_x F(x) + \lambda \|x\|_1\) with smooth \(F\) is ISTA (the iterative shrinkage-thresholding algorithm), i.e. proximal gradient descent: the \(L^1\) term is not differentiable, but its proximal operator is the closed-form soft threshold, so each iteration is one gradient step on \(F\) followed by one shrinkage,

\[ x^{(t+1)} \;=\; \mathcal{S}_{\eta\lambda}\big(x^{(t)} - \eta \nabla F(x^{(t)})\big), \qquad \mathcal{S}_\tau(v) \;=\; \operatorname{sign}(v) \cdot \max(|v| - \tau,\, 0). \]

The shrinkage sets small coordinates exactly to zero — which is why the iterates are genuinely sparse rather than merely small, and why the method fits here: the zero pattern of \(S\) is what keeps the compressed representation sparse. For the bilinear objectives above, alternate: with \(G'\) fixed the \(S\)-subproblem is exactly a LASSO (ISTA applies, with the gradient a sparse tensor contraction); with \(S\) fixed the \(G'\)-subproblem is plain least squares. FISTA’s momentum variant accelerates the \(O(1/t)\) convergence to \(O(1/t^2)\) at no extra cost per step.

The degree-3 case#

Write the value as a symmetric cubic form \(G(\varepsilon, \varepsilon, \varepsilon)\) (lower degrees via the sentinel symbol). Compress by learning low-dimensional feature maps of each degree,

\[ t^{(1)} = S_1(\varepsilon), \qquad t^{(2)} = S_2(\varepsilon, \varepsilon), \qquad t^{(3)} = S_3(\varepsilon, \varepsilon, \varepsilon), \]

and reconstruction forms that combine them to total degree three. Since every \(S_d\) is linear in the monomial features, the reconstruction has an explicit coefficient tensor — e.g. the pure-linear branch contributes \(G'_3 \cdot (S_1 \otimes S_1 \otimes S_1)\), with entries \(\sum_{abc} G'_{3,abc}\, S_{1,ai} S_{1,bj} S_{1,ck}\) — so the objective can be posed directly as a tensor norm on coefficients:

\[ \min_{G', S}\; \Big\| G - G'_3 \cdot (S_1 \otimes S_1 \otimes S_1) - G'_2 \cdot (S_2 \otimes S_1) - G'_1 \cdot S_3 \Big\|_F^2 \;+\; \lambda \sum_{d} \|S_d\|_1 . \]

Why the gradients are cheap#

The objective is a quadratic function of the residual and multilinear in each unknown factor — every \(G'_d\) and every \(S_d\) appears linearly in the reconstruction. The partial derivative with respect to any one factor is therefore just another tensor contraction of the same operands, with that factor’s slot left open. One example carries the whole idea: keep only the pure-linear branch,

\[ F \;=\; \big\| R \big\|_F^2 , \qquad R_{ijk} \;=\; G_{ijk} - \hat G_{ijk} \;=\; G_{ijk} - \sum_{abc} G'_{3,abc}\, S_{1,ai}\, S_{1,bj}\, S_{1,ck} . \]

The split matters computationally: the two parts of \(R\) are contracted by different machinery, and \(R\) itself is never materialized. \(G_{ijk}\) is the value’s sparse coefficient tensor, so contractions against it are sparse kernels over nse rows; \(\hat G_{ijk}\) is defined by a small tensor network (\(G'_3\) and three copies of \(S_1\)), so contractions against it never form the dense \(i j k\) tensor — a contraction planner such as cuTensorNet picks the optimal intermediate order (e.g. absorbing one \(S_1\) at a time), keeping every intermediate at the small \(k'\) widths.

Differentiating the squared norm gives \(-2R\) times the derivative of the reconstruction, and since \(G'_3\) appears linearly,

\[ \frac{\partial F}{\partial G'_{3,abc}} \;=\; -2 \sum_{ijk} R_{ijk}\, S_{1,ai}\, S_{1,bj}\, S_{1,ck} , \]

— the residual contracted with three copies of \(S_1\): one einsum. For \(S_1\), which appears three times, the product rule yields three structurally identical terms (equal, once \(G'_3\) and \(R\) are kept symmetric):

\[ \frac{\partial F}{\partial S_{1,ai}} \;=\; -6 \sum_{jk} R_{ijk} \sum_{bc} G'_{3,abc}\, S_{1,bj}\, S_{1,ck} , \]

— again one contraction, and both split over \(R\)’s two parts: the \(G\)-part is a sparse einsum (TACO or a custom Triton kernel), the \(\hat G\)-part a dense-but-tiny tensor-network contraction (cuTensorNet). The gradient computation therefore reuses exactly the kernels that evaluate the objective and needs no autodiff tape.

Considerations#

  1. Certified residual (difficulty 1). The noise added to the compressed value must bound \(\|f - g\|_\infty\) soundly, and the sound bound is coefficient-wise: \(\|r\|_\infty \le \sum |r_{\mathrm{coeff}}|\). The difficulty is that the absolute value does not commute with contraction — evaluating \(\sum |G - \hat G|\) exactly means expanding the tensor network \(\hat G\) into the dense coefficient tensor that the whole optimization was designed never to form. Candidate ways around the expansion:

    • Blockwise expansion: materialize the \(ijk\) grid tile by tile (as the blockwise DeepT matmul bound does) and accumulate \(\sum |R|\); the dense cost is paid once at certification, not per optimizer iteration.

    • Cauchy–Schwarz: \(\sum_r |r| \le \sqrt{N}\, \|R\|_F\), and \(\|R\|_F^2\) is tensor-network computable — expanding the square into \(\langle G, G\rangle - 2\langle G, \hat G\rangle + \langle \hat G, \hat G\rangle\) leaves only contractions (it is the training objective itself) — at the price of the \(\sqrt{N}\) slack.

    • Triangle inequality through the factors: \(\sum |\hat G| \le \sum_{abc} |G'_{3,abc}| \, \|S_{1,a\cdot}\|_1 \|S_{1,b\cdot}\|_1 \|S_{1,c\cdot}\|_1\) — cheapest, but bounds \(\sum|G| + \sum|\hat G|\) rather than the difference, so all cancellation between \(G\) and \(\hat G\) is lost.

    Floating point in the optimizer is harmless — only this final certification step must be sound.

  2. Choosing the widths \(k'_d\) (difficulty 2). A budget/energy rule (smallest \(k'\) with certified residual below a fraction of the value’s width), spectral decay of an HOSVD initialization, or a fixed per-layer compute budget; needs experiments.

  3. Cross-value correlation. Fresh symbols \(\varepsilon'\) forget their dependence on \(\varepsilon\). Compressing one tensor in isolation destroys its correlation with every other live value (attention weights vs. the residual stream), exactly like to_intervals does today. The objective should therefore compress all live values of a layer jointly, sharing one \(S\), and likely keep the input symbols uncompressed.

  4. Nonconvexity. The objective is bilinear in \((G', S)\) — alternating least squares or Adam from an SVD/symmetric-CP initialization; local minima are acceptable because any factorization is sound after certification, only tightness varies.

3. Polynomial Approximation of Non-linear Functions#

In principle every activation can be transferred polynomially: for \(f\) and an input value \(p(\varepsilon)\), choose a univariate polynomial \(q \approx f\) on the range of \(p\) and take \(q(p(\varepsilon)) + \delta\) with \(|\delta| \le \|f - q\|_{\infty, [\,\mathrm{lb}(p),\, \mathrm{ub}(p)\,]}\) — a degree-\((\deg q \cdot \deg p)\) polynomial with a certified remainder, keeping far more correlation than an affine enclosure.

Experiments in polysp.legendre show the naive version is currently worse: composition multiplies the monomial count, the degree cap truncates most of the new terms straight into noise, and that truncation noise exceeds what the affine linearizer would have paid. The prerequisite is §2 — with compression keeping nse bounded, composition can retain enough of \(q \circ p\) for the correlation gain to win. To revisit afterwards:

  • degree selection per element (crossing vs. saturated inputs need different \(q\));

  • certifying the remainder on the propagated range of \(p\), which is itself an abstract quantity;

  • ordering: compress before or after composing, and how §1’s basis change interacts with both.

Milestones#

  1. Legendre concretization inside lbub() only; measure width vs. cost on the attention benchmarks.

  2. Zonotope-case LASSO compression with certified residual; validate soundness by sampling.

  3. Degree-3 compression (shared \(S\) per layer); then re-run the polynomial activation-transfer experiments on top.