Beyond Quadratic Attention: Separable Gaussian Operators for ODE Vision Transformers
Organizations: Universitat Aut`onoma de Barcelona Centre de Visi´o per Computador (CVC) · Universitat Aut`onoma de Barcelona Centre de Recerca Matem`atica (CRM)
Abstract
Self-attention carries a quadratic cost in token count, which limits how far Vision Transformers can scale with input resolution. We introduce SSoG, a separable sum-of-Gaussians spatial operator that replaces dense self-attention with a learned mixture of axis-factorized Gaussian kernels, reducing the per-layer cost from to subquadratic in . Unlike dot-product self-attention, whose Lipschitz constant is unbounded outside a compact input domain, SSoG with fixed kernels is globally Lipschitz continuous, which, via the Picard--Lindelöf theorem, guarantees that the continuous-depth variant has a unique solution on the whole integration interval, a global guarantee that dense-attention counterparts lack in general. On CIFAR-100 at , in a matched discrete-Transformer setting and in the continuous setting, SSoG improves Top-1 accuracy by 2.2--4.1 points across three model widths while using 8.8--16.6% fewer parameters, with depth and training recipe held fixed. In a looped-Transformer setting, SSoG improves Top-1 accuracy by 5.4--8.1 points. We further probe inference-time scaling by evaluating on inputs up to at matched batch size: in the continuous-depth (Neural ODE) architecture, SSoG reaches up to 154% higher throughput and over 60% lower peak memory than dense attention, with the gap widening as resolution grows, consistent with the difference in asymptotic cost. In the discrete and looped settings, SSoG also achieves over 30% higher throughput than dense attention at this resolution. Together, these results establish SSoG as a practical, parameter-efficient, and theoretically grounded alternative to dense attention across discrete, weight-shared, and continuous-depth Transformer variants.
Figures & tables
| Variant | Acc@1 | Acc@3 | Acc@5 | Loss | Time (h) | |||
|---|---|---|---|---|---|---|---|---|
| Fixed | - | - | - | 74.22 | 88.59 | 92.38 | 1.12 | 47.77 |
| Full, cold | 70.89 | 86.89 | 91.30 | 1.20 | 59.48 | |||
| Center only | - | - | 73.33 | 87.79 | 91.61 | 1.15 | 58.37 | |
| Width only | - | - | 74.01 | 88.39 | 92.16 | 1.11 | 57.42 | |
| Mixture only | - | - | 73.87 | 88.20 | 91.73 | 1.13 | 57.08 | |
| Center + width | - | 71.21 | 86.83 | 91.13 | 1.19 | 58.61 |
| Continuous | Discrete | Looped | ||||
|---|---|---|---|---|---|---|
| Metric | MHA | SSoG | MHA | SSoG | MHA | SSoG |
| Acc@1 | 63.04 | 64.02 | 76.47 | 77.89 | 73.72 | 73.37 |
| Acc@3 | 79.04 | 79.62 | 89.44 | 90.59 | 88.35 | 87.93 |
| Acc@5 | 83.40 | 84.45 | 92.31 | 93.37 | 91.90 | 91.65 |
Appendix figures & tables15 assets
Supplementary material from the paper’s appendix.
Appendix
| Resolution | Tokens ( ) | SSoG | Dense (fused MHA) | Ratio ( ) |
|---|---|---|---|---|
| Whole model | Attention | ||||||
|---|---|---|---|---|---|---|---|
| Resolution | Tokens | MHA | SSoG | Reduction (%) | MHA | SSoG | Reduction (%) |
| 28.4 | 6.0 | ||||||
| 84.3 | 18.4 | ||||||
| 151.1 | 33.9 | ||||||
| 237.9 | 54.9 | ||||||
| 345.3 | 81.7 | ||||||
| Sink score | |||||||
| Setting | Operator | Mean | Worst layer | Rollout | Same key (%) | CLS (%) | High-norm (%) |
| Continuous | MHA | 11.3 | 22.9 | 1.7 | 6 | 0.3 | 0.0 |
| SSoG | 2.2 | 2.2 | 1.3 | 100 | - | 0.2 | |
| Discrete | MHA | 3.8 | 8.1 | 5.0 | 49 | 0.9 | 0.0 |
| SSoG | 2.6 | 3.3 | 1.5 | 100 | - | 0.3 | |
| SSoG steered | 3.3 | 5.3 | 1.7 | 60 | - | 0.7 | |
| Acc@1 | Val. loss | Learned min | |
|---|---|---|---|
| 0.10 | 73.73 | 1.14 | 0.182 |
| 0.25 | 74.22 | 1.12 | 0.366 |
| 0.50 | 73.02 | 1.14 | 0.578 |
| Setting | Variant | Steps | Acc@1 | Acc@3 | Acc@5 | Time (h) |
|---|---|---|---|---|---|---|
| Continuous | Fixed | 6 | 69.70 | 85.74 | 90.23 | 25.5 |
| Fixed | 8 | 72.16 | 87.30 | 91.05 | 33.2 | |
| Fixed | 12 | 74.22 | 88.59 | 92.38 | 47.8 | |
| Steered | 6 | 70.36 | 86.04 | 90.67 | 31.5 | |
| Steered | 8 | 72.00 | 87.16 | 91.42 | 40.3 | |
| Steered | 12 | 70.89 | 86.89 | 91.30 | 59.5 |
| Initialization | Final Acc@1 | Val. loss | Instabilities |
|---|---|---|---|
| Cold, full steering | 70.89 | 1.2017 | None |
| Warm, full steering | 65.56 | 1.3714 | Early plateau |
| Center-only, | 73.33 | 1.1546 | None |
| Center-only, | 71.53 | 1.2138 | None |
| Acc@1 | Val. loss | Mean | |
|---|---|---|---|
| 2 | 72.84 | 1.1385 | 0.529 |
| 4 | 72.48 | 1.1643 | 0.535 |
| 7 | 70.89 | 1.2017 | 0.475 |
| Setting | MHA Params(M) | SSoG Params | MHA Acc@1 | SSoG Acc@1 | |
|---|---|---|---|---|---|
| Continuous | 2.8 | 71.37 | |||
| 5.9 | 73.48 | ||||
| 10.3 | 74.22 | ||||
| Discrete | 18.5 | 77.56 | |||
| 41.0 | 78.18 | ||||
| 71.6 | 78.63 |
| Family | Ablation | Depth | Par(M) | Acc@1 | Acc@3 | Acc@5 | |||
| Discrete | Heads | 768 | 6 | 4 | 12 | 71.6 | 79.97 | 89.99 | 91.94 |
| 12 | 78.63 | 90.12 | 93.38 | ||||||
| 24 | 79.52 | 88.59 | 90.14 | ||||||
| Atoms | 768 | 12 | 1 | 12 | 71.6 | 79.62 | 89.38 | 91.52 | |
| 2 | 80.34 | 89.83 | 91.69 | ||||||
| 4 | 78.63 | 90.12 | 93.38 |
| Family | Ablation | Depth | Params | Acc@1 | Acc@3 | Acc@5 | ||
|---|---|---|---|---|---|---|---|---|
| Discrete | Heads | 768 | 6 | 12 | 85.9M | 75.63 | 87.51 | 90.55 |
| 12 | 74.67 | 86.89 | 90.19 | |||||
| 24 | 73.70 | 86.0 | 89.30 | |||||
| Depth | 768 | 12 | 6 | 43.3M | 70.47 | 84.50 | 88.61 | |
| 8 | 57.5M | 72.22 | 85.68 | 89.60 | ||||
| 12 | 85.9M | 74.67 | 86.89 | 90.19 |
| Patch 16 | Patch 14 | Patch 8 | ||||
|---|---|---|---|---|---|---|
| Setting | Operator | ( ) | ( ) | ( ) | ||
| Continuous | MHA | ( | ) | ( | ) | |
| SSoG | 74.20 | 74.85 ( | ) | ( | ) | |
| SSoG (rescaled ) | 74.20 | ( | ) | 64.39 ( | ) | |
| Discrete | MHA | ( | ) | ( | ) | |
| SSoG | 78.63 | ( | ) | ( | ) | |
| Acc@1 | Acc@5 | ||
|---|---|---|---|
| 384 | 6 | 55.65 (245) | 78.46 |
| 576 | 9 | 62.43 (273) | 83.37 |
| 768 | 12 | 64.02 (254) | 84.45 |
| Augmentation | Best Acc@1 (epoch) | Final Acc@1 | Runtime (h) |
|---|---|---|---|
| Strong | 74.22 (488) | 74.22 | 47.77 |
| Moderate | 73.09 (487) | 72.72 | 34.08 |
| Weak | 62.70 (452) | 62.27 | 33.26 |
| Setting | Protocol / Augmentation | Best Acc@1 | |
|---|---|---|---|
| Continuous | Linear probe (frozen) | Hard | 66.26 (43) |
| Linear probe (frozen) | Weak | 70.01 (45) | |
| Full fine-tune | Hard | 80.25 (91) | |
| Full fine-tune | Moderate | 83.84 (83) | |
| Discrete | Linear probe (frozen) | Hard | 66.77 (48) |
| Linear probe (frozen) | Weak | 69.05 (43) |
Transformer Attention
Transformer attention lets models relate tokens by weighting information from other positions, enabling context-dependent sequence and multimodal processing.
See all 333 papersContinuum Attention
Projection-Free Transformers via Gaussian Kernel Attention
nanochat framework, we implement causal masking and sliding-window constraints by masking and renormalizing the Gaussian kernel. At depth 20, a GKA model with the parameters and the total training FLOPs of a standard attention baseline trains stably, exhibits a near-zero train-validation gap, and demonstrates competitive behavior on standard benchmarks, albeit with higher bits-per-byte (BPB) at this compute scale. Overall, GKA provides a minimal, interpretable attention mechanism with an explicit locality scale, offering a dimension in the accuracy-efficiency trade-off for Transformer design.