Beyond Quadratic Attention: Separable Gaussian Operators for ODE Vision Transformers

Carlos Boned Riera; David Romero Sanchez; Oriol Ramos Terrades

cs.LGNov 5, 2026

Beyond Quadratic Attention: Separable Gaussian Operators for ODE Vision Transformers

Authors: Carlos Boned Riera, David Romero Sanchez, Oriol Ramos Terrades

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 $\mathcal{O}(N^2)$ to subquadratic in $N$ . 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 $224\times224$ , 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 $1024\times1024$ 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

Figure 1: Left: throughput–memory trade-off at resolution 1024 as the width of the model increases. Points farther upward provide more throughput and points farther left allocate less memory. Our proposed method dominates multi-head attention, achieving higher throughput at lower memory. Right: comparison of the three model families with our proposed method.

Figure 2: Width scaling on CIFAR-100, dense attention (MHA) versus SSoG, across the three architecture settings. The dotted connector and label give the Acc@1 gain of SSoG at $D=768$ .

Figure 3: Ablations of SSoG on CIFAR-100 in the three architecture settings.

Variant	$\mu$	$\sigma$	$\lambda$	Acc@1	Acc@3	Acc@5	Loss	Time (h)
Fixed	-	-	-	74.22	88.59	92.38	1.12	47.77
Full, cold	$\checkmark$	$\checkmark$	$\checkmark$	70.89	86.89	91.30	1.20	59.48
Center only	$\checkmark$	-	-	73.33	87.79	91.61	1.15	58.37
Width only	-	$\checkmark$	-	74.01	88.39	92.16	1.11	57.42
Mixture only	-	-	$\checkmark$	73.87	88.20	91.73	1.13	57.08
Center + width	$\checkmark$	$\checkmark$	-	71.21	86.83	91.13	1.19	58.61

Table 1: Steering-component attribution on CIFAR-100 at NFE 12 and

R=4

, cold initialization throughout. Checkmarks denote query-dependent components.

Figure 4: SSoG versus dense attention, throughput and memory improvement (%) by resolution, for each of the three architecture settings (continuous-depth ODE-ViT, discrete ViT-B/16, looped bViT). Each point is the mean relative difference over batch sizes $8\to 512$ at that resolution. Memory is negated (plotted as % memory saved ) so both lines share one direction: above the dashed zero line means SSoG wins that metric at that resolution, below means dense wins. The rightmost panel shows the relative improvement in FLOPs for the whole model and for the attention operation alone.

	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

Table 2: ImageNet-1k accuracy at

D=768

: SSoG versus dense attention (MHA).

Appendix figures & tables15 assets

Supplementary material from the paper’s appendix.

Appendix

Resolution	Tokens ( $N$ )	SSoG	Dense (fused MHA)	Ratio ( $\times$ )
$224$	$197$	$47.6$	$29.2$	$1.63$
$384$	$577$	$113.0$	$56.0$	$2.02$
$512$	$1025$	$184.5$	$88.3$	$2.09$
$640$	$1601$	$281.3$	$128.1$	$2.20$
$768$	$2305$	$395.8$	$181.4$	$2.18$
$1024$	$4097$	$689.0$	$304.3$	$2.26$

Appendix

Table 3: Peak allocated memory (MiB) of the attention operator alone, at matching token count.

		Whole model			Attention
Resolution	Tokens	MHA	SSoG	Reduction (%)	MHA	SSoG	Reduction (%)
$224$	$197$	$35.1$	28.4	$19.2$	$12.6$	6.0	$52.6$
$384$	$577$	$111.0$	84.3	$24.0$	$44.9$	18.4	$59.1$
$512$	$1025$	$214.1$	151.1	$29.4$	$96.8$	33.9	$65.0$
$640$	$1601$	$368.3$	237.9	$35.4$	$185.1$	54.9	$70.4$
$768$	$2305$	$590.1$	345.3	$41.5$	$326.4$	81.7	$75.0$

Appendix

Table 4: Forward FLOPs per image (GFLOPs, 2 per multiply-add) at dim768 with 12 block applications. The counts are identical, to within 0.3 GFLOPs, for the continuous (12 ODE steps), discrete (12 blocks) and looped (12 iterations) settings. Attention includes the projections; element-wise operations are not counted for either operator. Reduction is SSoG relative to dense SDPA attention (MHA).

		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
Continuous	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

Appendix

Table 5: Attention-sink diagnostics of trained models on 512 validation images. First 3 settings are on CIFAR100 and the last setting on ImageNet. Sink score : attention received by the most-attended key relative to uniform (

1

= no sink), averaged over heads, as the mean over layers, for the worst layer, and for the residual attention rollout

\prod_{l}(A_{l}+I)/2

of the head-averaged maps. Same key : share of images whose most-attended key is the head’s most frequent one (averaged over heads; most consistent layer); for fixed SSoG it is

100\%

by construction, because its content-independent kernel always peaks at an image corner through border renormalization. CLS : attention mass from patch queries to the class token (- where the model has none; SSoG patches never attend to it). High-norm : share of final pre-norm patch tokens with norm above

3\times

the per-image median.

$\sigma_{\min}$	Acc@1	Val. loss	Learned $\sigma$ 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

Appendix

Table 6: Effect of the sigma floor

\sigma_{\min}

on fixed SSoG.

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

Appendix

Table 7: Fixed versus steered SSoG on CIFAR-100 across the three architecture settings. Steps are ODE function evaluations (continuous), blocks (discrete) or loop iterations (looped). Time is the W&B training runtime for 500 epochs.

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, $\tilde{\alpha}_{\mu}=-8$	73.33	1.1546	None
Center-only, $\tilde{\alpha}_{\mu}=-4$	71.53	1.2138	None

Appendix

Table 8: Effect of steering initialization at NFE 12,

R=4

$M$	Acc@1	Val. loss	Mean $\|\mu\|$
2	72.84	1.1385	0.529
4	72.48	1.1643	0.535
7	70.89	1.2017	0.475

Appendix

Table 9: Effect of the steering radius

M

, full steering, NFE 12.

Setting	$D$	MHA Params(M)	SSoG Params	MHA Acc@1	SSoG Acc@1
Continuous	$384$	$3.07$	2.8	$69.15$	71.37
	$576$	$6.59$	5.9	$69.36$	73.48
	$768$	$11.45$	10.3	$70.47$	74.22
Discrete	$384$	$21.69$	18.5	$74.51$	77.56
	$576$	$48.47$	41.0	$74.20$	78.18
	$768$	$85.86$	71.6	$74.67$	78.63

Appendix

Table 10: Width scaling on CIFAR-100, dense attention versus SSoG, across the three architecture settings. Params are per model; bold marks the smaller parameter count and the better Acc@1 within each row.

Family	Ablation	$D$	$H$	$R$	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

Appendix

Table 11: Capacity and depth ablations on CIFAR-100 for the three SSoG model families. The Depth column gives the number of blocks for the Discrete family, the number of loop iterations for the Looped family, and the number of ODE function evaluations (NFE) for the Continuous family. The Heads and Atoms blocks are evaluated at NFE 12 where applicable, while the Depth block varies depth or NFE at fixed

D=768

H=12

R=4

. The column highlighted in blue is the parameter varied within each block; bold marks the best metric within each block.

Family	Ablation	$D$	$H$	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

Appendix

Table 12: MHA capacity and depth ablations on CIFAR-100 for the Discrete MHA (ViT) and Looped MHA (bViT) families. The column highlighted in blue is the parameter varied within each block; bold marks the best metric within each block.

Setting	Operator	( $14\times 14$ )	( $16\times 16$ )		( $28\times 28$ )
		Patch 16	Patch 14		Patch 8
Continuous	MHA	$70.47$	$70.23$ (	$+0.80$ )	$58.05$ (	$-11.38$ )
	SSoG	74.20	74.85 (	$+0.65$ )	$57.44$ (	$-16.76$ )
	SSoG (rescaled $\mu,\sigma$ )	74.20	$74.30$ (	$+0.10$ )	64.39 (	$-9.81$ )
Discrete	MHA	$74.67$	$74.91$ (	$-0.18$ )	$56.92$ (	$-18.17$ )
	SSoG	78.63	$79.34$ (	$+0.71$ )	$63.49$ (	$-15.53$ )

Appendix

Table 13: Patch-size robustness on CIFAR-100: models trained at patch 16 (224px,

14\times 14

tokens) evaluated at patch 14 and 8 without retraining. Patch-embedding kernels are PI-resized and position embeddings bicubically interpolated for every model; the rescaled SSoG row also scales

\mu

\sigma

by the grid ratio so each Gaussian covers the same image area. Acc@1 (%), change vs. patch 16 in parentheses.

$D$	$H$	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

Appendix

Table 14: ImageNet width sweep, fixed SSoG,

R=4

, NFE 12.

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

Appendix

Table 15: From-scratch augmentation ablation,

D=768

R=4

, NFE 12.

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)

Appendix

Table 16: ImageNet

\to

CIFAR-100 transfer,

D=768

backbone from an ImageNet checkpoint, across all three architecture settings.

Figure 5: Throughput and peak memory versus input resolution, dense attention versus SSoG, across all three architecture settings (rows: continuous-depth ODE-ViT, discrete ViT-B/16, looped bViT) and six batch sizes (columns: 8, 32, 64, 128, 256, 512), at the paper’s flagship width ( $D=768$ , $H=12$ ; NFE 12 for the continuous-depth setting, depth 12 for the discrete and looped settings). Each panel shares a log-scaled resolution axis (224-1024) with two y-axes: throughput in images/s (blue, left) and peak allocated memory in MiB (magenta, right). Solid circles/lines denote dense attention; dashed triangles denote SSoG. An $\times$ (dense) or $+$ (SSoG) marks the first resolution at which that configuration ran out of GPU memory; its vertical position is the last successfully measured value at that batch size, not a measured point at the marked resolution, and no further resolutions are plotted for that curve beyond it. All measurements are FP16 autocast (model parameters and inputs in FP32) on a single NVIDIA A40, randomly initialized weights, a single run per configuration (no repeated-seed or repeated-measurement uncertainty), 10 warmup and 50 measured iterations per point.

Transformer Attention

Transformer attention lets models relate tokens by weighting information from other positions, enabling context-dependent sequence and multimodal processing.

See all 333 papers

Self-Attention And Beyond the Infinite: Towards Linear Transformers with Infinite Self-Attention
Feb 26, 2026Giorgio Roffo, Hazem Abdelkawy, Nilli Lavie +1Self-Attention Efficient ViTs
Equixly API Security, Italy · Toyota Motor Europe, Belgium · MindVisionLabs, UK +2
Continuum Attention
Oct 14, 2026Afshin KhadangiCausal Attention Transformer Attention
SnT, University of Luxembourg
Projection-Free Transformers via Gaussian Kernel Attention
May 4, 2026Debarshi Kundu, Archisman Ghosh, Swaroop Ghosh +1Transformer Self-Attention
The Pennsylvania State University, University Park, USA

Feb 26, 2026J/K move · Enter open · S save

Self-Attention And Beyond the Infinite: Towards Linear Transformers with Infinite Self-Attention

Giorgio Roffo, Hazem Abdelkawy, Nilli Lavie, Luke Palmer

Equixly API Security, Italy · Toyota Motor Europe, Belgium · MindVisionLabs, UK+2

Self-Attention Efficient ViTs Vision Transformer Linear Attention

Open paper AskPDF arXiv

The quadratic cost of softmax attention limits Transformer scalability in high-resolution vision. We introduce Infinite Self-Attention (InfSA), a spectral reformulation that treats each attention layer as a diffusion step on a content-adaptive token graph, accumulating multi-hop interactions through a discounted Neumann series over attention matrices. This links self-attention to classical graph centrality (Katz, PageRank, eigenvector centrality) for interpretable token weighting. We also show the Neumann kernel equals the fundamental matrix of an absorbing Markov chain, so a token's centrality is its expected number of random-walk visits before absorption. We then propose Linear-InfSA, a linear-time variant that approximates the principal eigenvector of the implicit attention operator without forming the full attention matrix. It keeps an auxiliary state of fixed size proportional to per-head dimension dh (independent of sequence length N), is drop-in compatible with Vision Transformers, and supports stable training at 4096 by 4096 and inference at 9216 by 9216 (about 332k tokens). In a 4-layer ViT (53.5M parameters, 59 GFLOPs at 224 by 224), Linear-InfSA reaches 84.7% top-1 on ImageNet-1K, a +3.2 point architectural gain over an equal-depth softmax ViT trained with the same recipe. On ImageNet-V2, InfViT variants outperform all compared baselines (up to 79.8% vs 76.8%), indicating robustness under distribution shift. On an A100 40GB GPU, Linear-InfViT runs at 231 images/s and 0.87 J/image (13x better throughput and energy than equal-depth ViT) and is the only tested model to complete 9216 by 9216 inference without out-of-memory. The linear approximation closely matches the dominant eigenvector of the quadratic operator (cosine 0.985).

Abstract

Transformer Attention

Explore similar work

Self-Attention And Beyond the Infinite: Towards Linear Transformers with Infinite Self-Attention

Continuum Attention

Projection-Free Transformers via Gaussian Kernel Attention