Skip to main content
Back to timeline
arXivSource publication:

Theoretical analysis of sparse attention: JS divergence is fixed by truncation mass alone, information loss grows with sequence length, and generalization bounds tighten with sparsity

Related research and updates

Synopsis

This work presents a unified theoretical analysis of sparse attention: it proves that the Jensen–Shannon divergence between full and sparse attention is determined solely by the discarded truncation mass with a closed-form expression, uses order statistics under i.i.d. sub-Gaussian attention-score assumptions to show that the high-probability lower bound on truncation mass increases with sequence length, and derives sparsity-tightened generalization bounds via Rademacher complexity and a Xu–Raginsky mutual-information bound with entropy and covering-number analysis, verifying the closed-form divergence and the growth of information loss on GPT-2 (124M) with WikiText-103.

Source-provided article image: On the Trade-off Between Information Loss and Generalization in Sparse Attention

(a) M = 𝒪 ⁡ ( L ) M=\mathcal{O}(\sqrt{L})

arXiv

Interpretation

The paper derives a closed-form expression for the Jensen–Shannon divergence between full and sparse attention distributions, showing it is a function of the truncation mass alone, strictly convex and strictly increasing on [0,1]. Prior theoretical discussion of sparse attention was largely graph-theoretic or empirical, lacking a quantitative reduction of approximation error to a single scalar; this work compresses information loss entirely into an interpretable truncation mass. Analytical derivation via a weight-scaling lemma and term-by-term KL computation, verified on GPT-2 (124M) attention rows from the WikiText-103 test split with top- truncation, where empirical and closed-form values agree within float32 precision.

Under i.i.d. sub-Gaussian attention scores with bounded projection spectral norms and embedding norms, the paper uses order statistics to bound the retained-to-discarded mass ratio with high probability, yielding a high-probability lower bound on truncation mass that increases with sequence length. It elevates the growth of information loss with sequence length from an empirical observation to a probabilistic quantitative statement, and shows the rate is governed by the sparsity-budget order and a coefficient jointly set by embedding scale, projection spectral norms, and head dimension, not by sparsity alone. Analytical proof combining sub-Gaussian tail bounds and order-statistics concentration; empirically, truncation mass rises with sequence length across six sparsity strategies, matching the direction of the theorem, though the bound is one to three orders of magnitude more conservative than measurements.

Using Dudley's integral and a volumetric argument, the paper derives a Rademacher complexity upper bound for M-sparse attention layers and a corresponding generalization bound under fixed sparsity patterns, showing that sparsity reduces hypothesis-class complexity and tightens the bound. This is a formal characterization of Rademacher complexity for sparse attention, extending the single-layer result to multiple layers via a Lipschitz constant and linking the generalization gap to the product of layer-wise Lipschitz constants and hypothesis capacity. Analytical derivation; experiments use Gaussian-noise randomization tests as an empirical proxy for Rademacher complexity, observing that the generalization gap widens as sparsity increases, consistent with the theoretical growth curves.

Combining the Xu–Raginsky mutual-information bound with an entropy decomposition and covering-number analysis of the sparse hypothesis class, the paper obtains a generalization bound depending on both sparsity and the discretization resolution of attention weights. It connects information-theoretic generalization analysis with the structural properties of sparse attention, using a combinatorial term for support-set complexity and a geometric term for sparse-weight resolution, complementing the JS-divergence view. Analytical derivation, explicitly restricted to fixed, input-independent support sets and sparse weights confined to an ε-net of the simplex.

Perspective

The analysis applies to fixed, input-independent sparse support settings, positioned by the authors as a baseline complexity bound that isolates the effect of sparsity itself, decoupled from additional complexity introduced by adaptive routing. For practitioners deploying Transformers on long sequences, the results give actionable guidance: the sparsity-budget order and a coefficient jointly determined by embedding scale, projection spectral norms, and head dimension govern how fast information loss accumulates, and at a fixed embedding dimension a smaller head dimension (equivalently more heads) slows the growth of truncation mass with sequence length, subject to the attention entropy-collapse constraint. The closed-form JS divergence allows direct estimation of approximation error given a truncation mass, while the entropy and covering-number bound brings weight discretization resolution into generalization assessment, suited to quantized implementations.

The theoretical bounds rest on assumptions of i.i.d. sub-Gaussian attention scores and bounded projection spectral norms and embedding norms; the authors also note that the i.i.d. sub-Gaussian model does not capture how strongly real attention concentrates on a few keys, so the bound is conservative rather than tight, differing from measurements by one to three orders of magnitude. Under data-dependent sparsity patterns such as Top-k or routing, the hypothesis class varies with input and standard Rademacher analysis does not directly apply, which the authors list as future work. The generalization experiments use Gaussian-noise randomization tests as an empirical proxy for Rademacher complexity, and its relation to the generalization gap under real data distributions remains to be further observed. In addition, several formulas and figure captions appear as placeholders in the loaded text, so specific constants and curve details cannot be directly checked from the text; readers needing exact values should consult the original figures and tables.

Sources