TANGO: Treating Tokens as Operators

Joshua Nunley

cs.LGAug 22, 2026

TANGO: Treating Tokens as Operators

Authors: Joshua Nunley

Organizations: Department of Informatics Luddy School of Informatics, Computing, and Engineering Cognitive Science Program Indiana University Bloomington

Abstract

Transformers separate cross-token mixing in self-attention from token-wise transformation in feed-forward networks. We ask whether combining these operations can lower predictive loss under fixed data and parameter budgets. To do so, we introduce the Token-Aggregated Nonlinear Gating Operator (TANGO) model. TANGO computes a nonlinear feature-wise gate at each source token. Attention averages these gates for each destination. The average modulates a linear projection of the destination and forms the diagonal core of a source-conditioned linear operator. We test this proposal by comparing full-prefix and windowed TANGO with looped and untied Transformers, the Gated Attention Unit (GAU), and Fast Linear Attention with a Single Head (FLASH) on web text, Lean formal mathematics, DeepMind Mathematics, and code. The comparison uses two parameter scales, two depths, and three seeds. Checkpoints are selected on development data and evaluated on held-out test data. At matched parameters and training data, full-prefix TANGO has the lowest mean test negative log-likelihood in all 16 settings. In eight additional combinations of size and dataset, its development loss never increases as depth rises from 4 to 8 to 16, whereas the looped Transformer's loss increases in four. Full-prefix TANGO is computationally expensive because it averages wide gates over every visible source. To reduce this cost, we evaluate a variant with three narrower gated-projection sets assigned to the first, middle, and last applications. Across four FineWeb-Edu settings, this variant achieves 3.26 to 3.45 times the throughput of TANGO and 75% to 96% that of the looped Transformer. Its mean development negative log-likelihood is lower than TANGO's in three settings and 0.023 higher in the fourth, while remaining lower than both Transformer baselines in all four.

Figures & tables

Dataset	Tier	Depth	TANGO	Windowed TANGO	Loop Transformer++	Untied Transformer++	Full-attention GAU	FLASH
FineWeb-Edu	70M	4	3.270	3.323	3.363	3.291	3.449	3.431
FineWeb-Edu	70M	8	3.189	3.249	3.296	3.250	3.290	3.309
FineWeb-Edu	160M	4	2.950	2.997	3.012	2.964	3.032	3.067
FineWeb-Edu	160M	8	2.847	2.897	2.907	2.905	2.907	2.962
Lean	70M	4	1.413	1.875	2.143	1.830	1.902	1.929
Lean	70M	8	1.319	1.865	1.992	1.546	1.926	1.895

Table 1: Test NLL at the checkpoint selected on development data, for models matched in parameters and training data. Each entry is a mean over three seeds, and lower is better. Bold marks the lowest mean in each row. Tier is the model size, and depth is the number of block applications. Stack denotes The Stack v3 and DM Math denotes DeepMind Mathematics, where the official interpolation test is used. Lean gives equal weight to source files and proof completion. Values are comparable within a row but not across rows, and Table 12 gives standard deviations.

Tier	Depth	TANGO NLL	TANGO (3 sets) NLL	Loop NLL	Untied NLL	TANGO (3 sets) tok/s	vs. TANGO	vs. Loop	vs. Untied
70M	8	3.212 $\pm$ 0.016	3.183 $\pm$ 0.005	3.325 $\pm$ 0.007	3.277 $\pm$ 0.010	48.4k	3.26 $\times$	0.96 $\times$	0.75 $\times$
70M	16	3.181 $\pm$ 0.010	3.176 $\pm$ 0.009	3.277 $\pm$ 0.001	3.296 $\pm$ 0.010	26.6k	3.38 $\times$	0.75 $\times$	0.47 $\times$
160M	8	2.870 $\pm$ 0.006	2.857 $\pm$ 0.010	2.935 $\pm$ 0.002	2.930 $\pm$ 0.003	27.6k	3.40 $\times$	0.91 $\times$	0.66 $\times$
160M	16	2.815 $\pm$ 0.013	2.838 $\pm$ 0.007	2.901 $\pm$ 0.005	2.937 $\pm$ 0.008	14.6k	3.45 $\times$	0.79 $\times$	0.37 $\times$

Table 2: TANGO with three gated-projection sets on FineWeb-Edu. NLL is the mean

\pm

standard deviation over three seeds on development data at the selected checkpoint, so these values differ from the test values in Table 1 . Throughput is the mean training rate at the end of training, and the final columns divide the three-set model’s throughput by that of each comparison model.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Tier	Depth	$d$	TANGO / Windowed	Loop	Untied	GAU	FLASH	Non-embedding range
70M	4	512	28,480	28,144	6,528	7,136–7,152	7,152	44.238M–44.307M
70M	8	512	28,480	28,144	2,912	3,536–3,552	3,552	44.180M–44.279M
160M	4	768	52,176	51,672	12,144	13,104–13,120	13,120	121.293M–121.429M
160M	8	768	52,176	51,672	5,568	6,528–6,544	6,544	121.272M–121.543M

Appendix

Table 3: Parameter-matched configurations for the primary comparison. Tier is the model size and depth is the number of block applications. Entries are gated, feed-forward, or GAU expansion widths, and GAU is given as a range because its learned relative-position parameters depend on context length. The Loop and Untied columns refer to the corresponding Transformer++ baselines.

Dataset	Tier	TANGO	Windowed TANGO	Loop Transformer++	Untied Transformer++	Conventional Untied	Full-attention GAU	FLASH
FineWeb-Edu	70M	3.181 $\pm$ 0.010	3.259 $\pm$ 0.008	3.277 $\pm$ 0.001	3.296 $\pm$ 0.010	3.293 $\pm$ 0.017	3.223 $\pm$ 0.004	3.266 $\pm$ 0.009
FineWeb-Edu	160M	2.815 $\pm$ 0.013	2.888 $\pm$ 0.005	2.901 $\pm$ 0.005	2.937 $\pm$ 0.008	2.936 $\pm$ 0.006	2.857 $\pm$ 0.007	2.904 $\pm$ 0.004
Lean	70M	1.477 $\pm$ 0.017	2.062 $\pm$ 0.028	2.247 $\pm$ 0.048	1.718 $\pm$ 0.144	1.619 $\pm$ 0.041	2.178 $\pm$ 0.023	2.179 $\pm$ 0.019
Lean	160M	1.570 $\pm$ 0.015	2.086 $\pm$ 0.018	2.333 $\pm$ 0.037	1.624 $\pm$ 0.013	1.592 $\pm$ 0.012	2.176 $\pm$ 0.026	2.213 $\pm$ 0.027
DM Math	70M	0.379 $\pm$ 0.006	0.384 $\pm$ 0.004	0.671 $\pm$ 0.144	0.511 $\pm$ 0.010	0.503 $\pm$ 0.005	0.448 $\pm$ 0.004	0.429 $\pm$ 0.008
DM Math	160M	0.240 $\pm$ 0.002	0.232 $\pm$ 0.004	0.753 $\pm$ 0.149	0.786 $\pm$ 0.012	0.849 $\pm$ 0.027	0.318 $\pm$ 0.003	0.314 $\pm$ 0.004

Appendix

Table 4: Best development NLL at depth 16, mean

\pm

standard deviation over three seeds. Bold marks the lowest mean in each row. This appendix reports development results only.

Dataset	Tier	Depth 4	Depth 8	Depth 16
FineWeb-Edu	70M	3.293	3.212	3.181 $\pm$ 0.010
FineWeb-Edu	160M	2.973	2.870	2.815 $\pm$ 0.013
Lean	70M	1.649	1.522	1.477 $\pm$ 0.017
Lean	160M	1.698	1.570	1.570 $\pm$ 0.015
DM Math	70M	0.473	0.386	0.379 $\pm$ 0.006
DM Math	160M	0.383	0.270	0.240 $\pm$ 0.002

Appendix

Table 5: Best development NLL of TANGO at depths 4, 8, and 16. Depth-4 and depth-8 values are the three-seed means of Table 14 , and depth-16 values are means

\pm

standard deviations over three seeds.

Tier	Depth	$d$	Heads	$F$	$F/d$	Non-embedding	Total
70M	4	800	25	2,064	2.580	30,061,600	70,267,200
70M	8	640	20	1,616	2.525	37,939,840	70,104,320
70M	16	480	15	1,344	2.800	45,727,200	69,850,560
160M	4	1,376	43	3,664	2.663	90,806,368	159,960,000
160M	8	1,056	33	2,816	2.667	107,071,008	160,142,400
160M	16	800	25	2,080	2.600	120,858,400	161,064,000

Appendix

Table 6: Geometry of the conventionally proportioned Untied baselines at depths 4, 8, and 16. The head dimension is 32. Counts use the GPT-2 vocabulary and count the tied input and output embedding once.

Dataset	Tier	Depth	TANGO	Original-grid Untied	Conventional Untied	$\Delta$ (conventional–original)
FineWeb-Edu	70M	4	3.293 $\pm$ 0.005	3.318 $\pm$ 0.006	3.321 $\pm$ 0.003	+0.003
FineWeb-Edu	70M	8	3.212 $\pm$ 0.016	3.277 $\pm$ 0.010	3.285 $\pm$ 0.003	+0.007
FineWeb-Edu	70M	16	3.181 $\pm$ 0.010	3.296 $\pm$ 0.010	3.293 $\pm$ 0.017	-0.003
FineWeb-Edu	160M	4	2.973 $\pm$ 0.008	2.990 $\pm$ 0.003	3.010 $\pm$ 0.002	+0.019
FineWeb-Edu	160M	8	2.870 $\pm$ 0.006	2.930 $\pm$ 0.003	2.958 $\pm$ 0.008	+0.028
FineWeb-Edu	160M	16	2.815 $\pm$ 0.013	2.937 $\pm$ 0.008	2.936 $\pm$ 0.006	-0.001

Appendix

Table 7: Development NLL at the selected checkpoint in all 24 settings. Each entry is a mean

\pm

standard deviation over three seeds. Bold marks the lowest mean among the three architectures shown, and the final column subtracts the original Untied result from the conventional one. This comparison uses development results only.

Architecture	Active packed width	Dense-trained	Rotating sparse	All segments active	Narrow specialist	$\Delta$ (sparse–specialist)
TANGO	368	3.212 $\pm$ 0.016	3.397 $\pm$ 0.012	3.353 $\pm$ 0.010	3.431 $\pm$ 0.005	-0.034
Windowed TANGO	368	3.282 $\pm$ 0.001	3.488 $\pm$ 0.002	3.446 $\pm$ 0.002	3.507 $\pm$ 0.004	-0.019
Loop Transformer++	5,632	3.325 $\pm$ 0.007	3.449 $\pm$ 0.004	4.223 $\pm$ 0.045	3.501 $\pm$ 0.004	-0.052
Untied Transformer++	640	3.277 $\pm$ 0.010	3.562 $\pm$ 0.012	4.214 $\pm$ 0.131	3.460 $\pm$ 0.008	+0.102

Appendix

Table 8: FineWeb-Edu development NLL at 70M and depth 8 under rotating narrow-width training. The rotating-sparse and all-segments-active columns run the same stored checkpoint, and the narrow specialist is a separately trained model. The final column subtracts the specialist’s NLL from the rotating-sparse NLL, so negative values favor the rotating model. Packed widths are per attention head for the TANGO variants and total feed-forward width for the Transformer++ variants.

Target	TANGO					TANGO (3 sets)
$F/d$	$d$	Heads	$F$	Realized	Params	$d$	Heads	$F$	Realized	Params
0.5	1,280	40	640	0.500	70.066M	1,216	38	570	0.469	70.310M
1	1,248	39	1,170	0.938	70.219M	1,120	35	1,120	1.000	70.089M
2	1,184	37	2,257	1.906	70.327M	992	31	2,015	2.031	69.815M
4	1,088	34	4,046	3.719	70.255M	864	27	3,267	3.781	70.321M
8	928	29	7,656	8.250	69.677M	704	22	5,346	7.594	70.246M

Appendix

Table 9: Model dimensions for the active-ratio comparison at depth 4. Parameter counts include the tied token embedding once.

Target $F/d$	TANGO NLL	TANGO (3 sets) NLL	$\Delta$ NLL	TANGO tok/s	TANGO (3 sets) tok/s
0.5	3.744 $\pm$ 0.009	3.575 $\pm$ 0.009	-0.169 $\pm$ 0.011	82.2k	95.1k
1	3.591 $\pm$ 0.009	3.446 $\pm$ 0.009	-0.145 $\pm$ 0.018	80.4k	103.9k
2	3.473 $\pm$ 0.005	3.363 $\pm$ 0.009	-0.111 $\pm$ 0.014	72.3k	95.7k
4	3.384 $\pm$ 0.015	3.305 $\pm$ 0.006	-0.079 $\pm$ 0.010	69.7k	95.9k
8	3.315 $\pm$ 0.013	3.277 $\pm$ 0.007	-0.037 $\pm$ 0.016	59.3k	90.9k
16	3.281 $\pm$ 0.014	3.271 $\pm$ 0.013	-0.009 $\pm$ 0.020	47.5k	73.0k

Appendix

Table 10: Active-ratio comparison at 70M and depth 4. NLL entries and the paired differences between the three-set model and TANGO are means

\pm

standard deviations over three paired seeds, and negative differences favor the three-set model. Rates are means over seeds of each run’s median over its final 100 updates.

Presented tokens	Seed pair	TANGO	Windowed TANGO	Untied	$\Delta$ (TANGO–Untied)
3.20 billion	17/101	2.745	2.823	2.835	-0.090

Appendix

Table 11: FineWeb-Edu development NLL at 410M and depth 8 after the same number of training tokens. Every entry uses seed pair

(17,101)

after 3.20 billion tokens.

Dataset	Tier	Depth	TANGO	Windowed TANGO	Loop Transformer++	Untied Transformer++	Full-attention GAU	FLASH
FineWeb-Edu	70M	4	3.270 $\pm$ 0.005	3.323 $\pm$ 0.005	3.363 $\pm$ 0.003	3.291 $\pm$ 0.006	3.449 $\pm$ 0.004	3.431 $\pm$ 0.020
FineWeb-Edu	70M	8	3.189 $\pm$ 0.016	3.249 $\pm$ 0.001	3.296 $\pm$ 0.006	3.250 $\pm$ 0.011	3.290 $\pm$ 0.019	3.309 $\pm$ 0.006
FineWeb-Edu	160M	4	2.950 $\pm$ 0.008	2.997 $\pm$ 0.002	3.012 $\pm$ 0.007	2.964 $\pm$ 0.003	3.032 $\pm$ 0.013	3.067 $\pm$ 0.019
FineWeb-Edu	160M	8	2.847 $\pm$ 0.006	2.897 $\pm$ 0.004	2.907 $\pm$ 0.001	2.905 $\pm$ 0.004	2.907 $\pm$ 0.012	2.962 $\pm$ 0.015
Lean	70M	4	1.413 $\pm$ 0.011	1.875 $\pm$ 0.016	2.143 $\pm$ 0.030	1.830 $\pm$ 0.027	1.902 $\pm$ 0.051	1.929 $\pm$ 0.046
Lean	70M	8	1.319 $\pm$ 0.021	1.865 $\pm$ 0.013	1.992 $\pm$ 0.021	1.546 $\pm$ 0.044	1.926 $\pm$ 0.034	1.895 $\pm$ 0.032

Appendix

Table 12: Test NLL at the selected checkpoint, mean

\pm

standard deviation over three seeds. Mathematics uses the official interpolation test, and bold marks the lowest mean in each row.

Dataset	Tier	Depth	Seed 17/101	Seed 23/103	Seed 41/107
FineWeb-Edu	70M	4	-0.0282	-0.0219	-0.0132
FineWeb-Edu	70M	8	-0.0338	-0.0761	-0.0588
FineWeb-Edu	160M	4	-0.0112	-0.0252	-0.0062
FineWeb-Edu	160M	8	-0.0444	-0.0608	-0.0455
Lean	70M	4	-0.4217	-0.4235	-0.3939
Lean	70M	8	-0.2409	-0.2840	-0.1566

Appendix

Table 13: Test NLL margin of TANGO over the best other model in the primary comparison, for each seed. Negative values favor TANGO. Mathematics uses the official interpolation test, and each value compares the six architectures at the same model and data-order seed.

Dataset	Tier	Depth	TANGO	Windowed TANGO	Loop Transformer++	Untied Transformer++	Full-attention GAU	FLASH
FineWeb-Edu	70M	4	3.293 $\pm$ 0.005	3.355 $\pm$ 0.005	3.392 $\pm$ 0.003	3.318 $\pm$ 0.006	3.484 $\pm$ 0.003	3.465 $\pm$ 0.020
FineWeb-Edu	70M	8	3.212 $\pm$ 0.016	3.282 $\pm$ 0.001	3.325 $\pm$ 0.007	3.277 $\pm$ 0.010	3.323 $\pm$ 0.018	3.342 $\pm$ 0.006
FineWeb-Edu	160M	4	2.973 $\pm$ 0.008	3.028 $\pm$ 0.002	3.041 $\pm$ 0.007	2.990 $\pm$ 0.003	3.061 $\pm$ 0.012	3.097 $\pm$ 0.019
FineWeb-Edu	160M	8	2.870 $\pm$ 0.006	2.930 $\pm$ 0.004	2.935 $\pm$ 0.002	2.930 $\pm$ 0.003	2.935 $\pm$ 0.011	2.992 $\pm$ 0.016
Lean	70M	4	1.649 $\pm$ 0.012	2.115 $\pm$ 0.029	2.441 $\pm$ 0.035	2.095 $\pm$ 0.034	2.202 $\pm$ 0.033	2.225 $\pm$ 0.047
Lean	70M	8	1.522 $\pm$ 0.017	2.103 $\pm$ 0.022	2.323 $\pm$ 0.014	1.784 $\pm$ 0.067	2.195 $\pm$ 0.028	2.175 $\pm$ 0.026

Appendix

Table 14: Development NLL at the checkpoint selected from each run’s evaluation schedule, mean

\pm

standard deviation over three seeds. These values cover all 288 training runs and were used to select checkpoints, so they are separate from the test results.

Dataset	TANGO	Low-dimensional variant	Paired $\Delta$	Throughput speedup / peak GPU-memory reduction
FineWeb-Edu	$3.2931\pm 0.0046$	$3.3668\pm 0.0044$	$0.0736\pm 0.0089$	$2.71\times$ / $8.2\%$
Stack	$0.8995\pm 0.0032$	$0.9619\pm 0.0055$	$0.0624\pm 0.0056$	not measured
Lean	$1.6489\pm 0.0116$	$1.8445\pm 0.0035$	$0.1955\pm 0.0097$	not measured
DeepMind Mathematics	$0.4733\pm 0.0081$	$0.5014\pm 0.0029$	$0.0281\pm 0.0109$	$1.92\times$ / $7.0\%$

Appendix

Table 15: Compact-controller TANGO at 70M and depth 4. Quality is development NLL at the checkpoint selected from the evaluation schedule, reported as the mean

\pm

standard deviation over three paired seeds, and the paired

\Delta

is compact-controller TANGO minus TANGO, so positive values favor TANGO. Systems measurements use matched complete optimizer updates with eager bfloat16 execution on prepared batches, on one H100 graphics processor with 80 gigabytes of memory. Speedup is compact-controller throughput divided by TANGO throughput, memory reduction refers to peak allocated graphics-processor memory, and “not measured” marks a task shape that was not measured.

Explore similar work

Apr 18, 2026cs.LG

Adaptive Computation Depth via Learned Token Routing in Transformers

Standard transformer architectures apply the same number of layers to every token regardless of contextual difficulty. We present Token-Selective Attention (TSA), a learned per-token gate on residual updates between consecutive transformer blocks. Each gate is a lightweight two-layer multi-layer perceptron (MLP) that produces a continuous halting probability, making the mechanism end-to-end differentiable with 1.7% parameter overhead and no changes to the base architecture. Notably, TSA learns difficulty-proportional routing without any explicit depth pressure: even at

λ=0

(no depth regularisation), the task-loss gradient alone drives the router to skip 20% of token-layer operations. On character-level language modeling, TSA saved 14-23% of token-layer operations (TLOps) across Tiny-Shakespeare and enwik8 at <0.5% quality loss. At matched efficiency, TSA achieved 0.7% lower validation loss than early exit, and the learned routing transfers directly to inference-time sparse execution for real wall-clock speedup.

Ahmed Abdelmuniem Abdalla Mohammed

Independent Researcher

Transformer Architectures Multilayer Perceptrons

Ask PDF

May 4, 2026cs.LG

Projection-Free Transformers via Gaussian Kernel Attention

Self-attention in Transformers is typically implemented as

\mathrm{softmax}(QK^\top/\sqrt{d})V

, where

Q=XW_Q

K=XW_K

, and

V=XW_V

are learned linear projections of the input

X

. We ask whether these learned projections are necessary, or whether they can be replaced by a simpler similarity-based diffusion operator. We introduce \textbf{Gaussian Kernel Attention} (GKA), a drop-in replacement for dot-product attention that computes token affinities directly using a Gaussian radial basis function (RBF) kernel applied to per-head token features. Each head learns only a bandwidth parameter

σ_h

, while a single output projection

W_O

preserves compatibility with the standard Transformer interface. GKA can be interpreted as normalized kernel regression over tokens, linking modern Transformer architectures to classical non-local filtering and kernel smoothing methods. We evaluate GKA in both vision and language modeling settings. For autoregressive language modeling within the \texttt{nanochat} framework, we implement causal masking and sliding-window constraints by masking and renormalizing the Gaussian kernel. At depth 20, a GKA model with

0.42\times

the parameters and

0.49\times

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.

Debarshi Kundu, Archisman Ghosh, Swaroop Ghosh +1

The Pennsylvania State University, University Park, USA

Transformer Attention Transformer Architectures Kernel Method Autoregressive Language Models

Ask PDF

May 21, 2026cs.LG

Energy-Gated Attention: Spectral Salience as an Inductive Bias for Transformer Attention

Standard transformer attention computes pairwise similarity between queries and keys, treating all tokens as equally salient regardless of their intrinsic informational content. In turbulent fluid dynamics, coherent structures -- the energetically dominant, spatially organized patterns that persist amid background chaos -- carry a disproportionate fraction of total energy and govern all transport. We propose that tokens play an analogous role in transformer attention: informationally dense positions (morphological boundaries, syntactic heads, discourse markers) concentrate spectral energy and should attract proportionally more attention than background tokens (function words, repeated patterns, low-information filler). We propose Energy-Gated Attention (EGA): a simple modification that gates value aggregation by the spectral energy of key token embeddings, computed by a single learned linear projection that discovers the dominant spectral mode of the embedding field. On TinyShakespeare, EGA achieves +0.103 validation loss improvement with only 12,480 additional parameters (<0.26% overhead) and no measurable computational cost. The result is consistent on Penn Treebank (+0.101), demonstrating dataset independence. A systematic ablation across three wavelet families (fixed Morlet, Daubechies db2/db4, and a parametric Morlet) establishes that fixed structured bases are suboptimal -- the optimal energy direction is data-adaptive and non-sinusoidal -- while identifying learned wavelet packets as a promising open direction. The learned energy threshold converges to tau ~= 0.35 independently of initialization, corresponding to the fraction (~36%) of tokens carrying above-average spectral energy in English text, a stable linguistic property consistent with the fraction of content words in running English text.

Athanasios Zeris

Independent Researcher, Athens, Greece.

Transformer Attention Token Embeddings Spectral Inductive Bias

Ask PDF