PaperlayerView exact revision on arXiv

Source-linked AI summary

TANGO: Token-Aggregated Nonlinear Gating Operators for Natural and Formal Language Modeling

Joshua Nunley

arXiv:2608.22117v1cs.LGcs.CL

TL;DR

The paper asks whether cross-token interaction and position-wise nonlinear transformation can be replaced by a single gated update. It introduces TANGO and the linear-complexity WANGO variant, finding that TANGO achieves the best validation NLL across the tested benchmarks, while WANGO leads among linear-complexity models. These results are limited to validation NLL and analytical operation counts at the tested scale and budgets.

  • Problem

    Standard Transformer blocks separate cross-token self-attention from a nonlinear feed-forward network applied independently at each position.

    Source: §1 · p. 1

  • Method

    TANGO replaces both sublayers with a cross-token gated residual update, while WANGO combines local pairwise weighting with positive-feature-map prefix statistics for older sources.

    Source: §1 · p. 2, §7 · p. 11

  • Results

    TANGO obtains the lowest validation NLL on FineWeb-Edu, Lean, and DeepMind Mathematics; WANGO leads among linear-complexity architectures on FineWeb-Edu.

    Source: §1 · p. 2, §7 · p. 11

  • Takeaways & Limitations

    The proposed gated-update models provide a full-prefix model with the strongest tested validation NLL and a windowed model with linear sequence-length complexity.

    Source: §7 · p. 11

    The experiments do not match forward-pass operation counts or establish scaling laws, and they do not measure generated-answer accuracy, proof acceptance, wall-clock speed, or GPU memory.

    Source: §6.3 · p. 10

Abstract

from arXiv · show

A standard Transformer block separates cross-token interaction in self-attention from a nonlinear feed-forward network applied independently at each position. We introduce the TANGO model (Token-Aggregated Nonlinear Gating Operators), which replaces these two sublayers with one cross-token gated residual update. Each source token produces a SwiGLU gate vector. Query-key similarities determine a weighted average of source gates for each destination, and the resulting gate rescales projected destination features. TANGO assigns a separate weight to every causally visible source and is quadratic in sequence length. The WANGO model (Windowed Aggregation of Nonlinear Gating Operators) retains the same unnormalized scores within a recent window and uses positive feature-map prefix statistics for older sources, giving linear sequence-length complexity for fixed window and feature dimensions. We compare TANGO and WANGO with Recurrent and Untied Transformer++, full-attention GAU, and FLASH. All models have approximately 44.3M nonembedding parameters and are trained in three matched runs. TANGO, WANGO, and Recurrent Transformer++ apply one shared block four times; the other architectures use four independent blocks. TANGO obtains the lowest mean validation negative log-likelihood on FineWeb-Edu, Lean, and DeepMind Mathematics, although it has the largest analytical forward-pass operation count. WANGO obtains the lowest mean FineWeb-Edu NLL among the architectures with computation linear in sequence length and outperforms Recurrent Transformer++ at nearly the same analytical forward-pass multiply-accumulate count.

1 Introduction

TANGO replaces self-attention and the position-wise feed-forward network with one cross-token gated residual update, while WANGO provides a linear-sequence-length variant. In matched comparisons, TANGO has the best validation NLL overall, and WANGO leads among linear-complexity architectures.

TANGO

  • Each destination averages source gates using query–key similarities, then uses the result to rescale its projected features.The gate is generated from the source representation and acts as a diagonal scaling operator on destination features.Source: §1 · pp. 1–2

WANGO

  • TANGO aggregates every causally visible source, whereas WANGO uses direct pairwise weighting in a recent window and positive-feature-map running sums for older sources.WANGO retains the same exponential similarity scores locally and factorizes older-source scores to maintain prefix statistics.Source: §1 · p. 2

Results

  • TANGO obtains the lowest mean validation NLL on FineWeb-Edu, Lean, and DeepMind Mathematics, while WANGO leads among linear-complexity architectures on FineWeb-Edu.WANGO also outperforms Recurrent Transformer++ at a nearly identical analytical forward-pass multiply–accumulate count.Source: §1 · p. 2

TANGO

  • TANGO replaces self-attention and the position-wise feed-forward network with one cross-token gated residual update.Source tokens produce SwiGLU gates, while destination positions apply weighted averages of those gates to their projected features.Source: §1 · pp. 1–2

WANGO

  • WANGO has computation linear in sequence length for fixed window and feature dimensions.The proposed models are compared with recurrent, untied, full-attention, and hybrid linear-attention architectures under matched experimental conditions.Source: §1 · p. 2

2 Model

TANGO and WANGO replace separate self-attention and position-wise feed-forward sublayers with source-conditioned gated residual updates. TANGO aggregates all causal sources quadratically, while WANGO combines recent-window scores with factorized older-prefix statistics for linear complexity.

Source-conditioned gating

  • TANGO and WANGO replace causal self-attention and position-wise feed-forward sublayers with one source-conditioned gated update.A single shared block is applied four times in the proposed models.Source: §2.1 · p. 3, §2.3 · p. 4
  • Each destination averages causally visible source gates using content-dependent weights, then applies the result elementwise to its projected features.The gate and projected features can originate from different positions; source gates act as diagonal scaling operators rather than arbitrary projection matrices.Source: §2.1 · p. 3

TANGO aggregation

  • TANGO uses temperature-scaled cosine similarities over every causal source, producing quadratic sequence-length complexity.Its query–key scoring costs Θ(T^2d), while gate aggregation costs Θ(T^2F).Source: §2.2 · pp. 3–4

WANGO aggregation

  • WANGO preserves exponential cosine-similarity scores in a trailing window and uses a positive feature map with running prefix sums for older sources.Recent and older contributions are normalized together, and every older source remains able to contribute.Source: §2.2 · p. 4
  • For fixed window and feature dimensions, WANGO has linear sequence-length computation while retaining the same parameterization and parameter count as TANGO.The experiments use W = 64 and compute prefix sums with blockwise parallel scans.Source: §2.2 · p. 4

Comparison architectures

  • The comparison includes Recurrent Transformer++, Untied Transformer++, full-attention GAU, and FLASH under approximately 44.3M nonembedding parameters.Recurrent Transformer++ shares one block across four applications; Untied Transformer++, GAU, and FLASH use independent blocks.Source: §2.3 · pp. 4–5, Table 1 · p. 5

3 Experimental Design

The study compares six approximately parameter-matched architectures across three benchmarks using matched seeds, data order, training budgets, and final-checkpoint validation NLL. Experiments include four-block and eight-block settings, with shared or independent block parameters depending on architecture.

Models and parameter matching

  • Six architectures target approximately 44.3M nonembedding parameters and perform four sequential block applications.TANGO, WANGO, and Recurrent Transformer++ reuse one block; Untied Transformer++, GAU, and FLASH use four independently parameterized blocks.Source: §3 · p. 5

Training and evaluation

  • Each architecture is trained three times with matched initialization seeds, data-order seeds, examples, training budgets, and teacher-forced evaluation.Validation NLL is reported as the arithmetic mean and sample standard deviation across final checkpoints.Source: §3 · pp. 5–6

Additional experiment

  • An eight-block comparison repeats the six architectures and three seed pairs on FineWeb-Edu and Lean under the same data splits, contexts, optimization, and evaluation procedure.The experiment contains 36 runs.Source: Eight-block comparison · p. 6

Benchmarks

  • FineWeb-Edu uses 8,192-token sequences from temporally partitioned documents and reports validation NLL per GPT-2 token.Training uses 5,722 optimizer steps and an effective batch size of 32 sequences.Source: FineWeb-Edu · p. 6
  • Lean combines causal Mathlib source-code modeling with proof completion, alternates task-specific batches, and gives both tasks equal optimization weight.Joint validation NLL is the arithmetic mean of source-code and proof-completion NLLs.Source: FineWeb-Edu · p. 6
  • DeepMind Mathematics samples equally from 168 module–difficulty combinations, supervises answer characters and EOS, and uses context length 256.The benchmark includes all 56 modules at easy, medium, and hard difficulty.Source: FineWeb-Edu · p. 6

Compute accounting

  • The study reports analytical MAC counts from model equations while excluding embedding lookup, normalization, masking, feature-map construction, and elementwise operations.These counts describe arithmetic work rather than measured runtime or GPU memory.Source: FineWeb-Edu · p. 7

4 Results

TANGO achieves the lowest validation NLL across the reported benchmarks and remains strongest with eight block applications. WANGO is the best linear-scaling architecture and improves on Recurrent Transformer++ at nearly identical analytical MAC counts.

Four-block results

  • TANGO has the lowest validation NLL on all three benchmarks, while WANGO has the lowest FineWeb-Edu NLL among linear-scaling architectures.WANGO also beats Recurrent Transformer++ at a similar analytical forward-pass operation count.Source: §4 · p. 7

Eight-block results

  • With eight block applications, TANGO has the lowest NLL on both FineWeb-Edu and Lean, while WANGO remains below Recurrent Transformer++ and FLASH on FineWeb-Edu.The comparison covers six models, two datasets, and three seeds.Source: §4.1 · p. 8, Table 4 · p. 8
  • 3.293 ± 0.005 is TANGO’s eight-block FineWeb-Edu NLL, while WANGO records 3.355 ± 0.005 versus 3.392 ± 0.003 for Recurrent Transformer++.FLASH records 3.465 ± 0.020 among the linear-scaling architectures.Source: §4.2 · p. 8

Compute comparison

  • 1.90T MACs for WANGO is close to 1.94T MACs for Recurrent Transformer++ on one 8,192-token sequence.For fixed window and feature dimensions, WANGO’s operation count grows linearly rather than quadratically with sequence length; wall-clock speed and GPU memory are not reported.Source: §4.2 · p. 8, Table 5 · p. 8

Lean and DeepMind Mathematics

  • 1.719 ± 0.026 is TANGO’s lowest joint Lean validation NLL, compared with 2.138 ± 0.033 for Untied Transformer++.The joint score equally averages source-code modeling and proof completion under teacher forcing.Source: §4 · p. 7, §4.3 · p. 8
  • 0.473 ± 0.008 is TANGO’s DeepMind Mathematics validation NLL, compared with 0.487 ± 0.006 for WANGO and 0.537 ± 0.003 for Recurrent Transformer++.NLL is measured per supervised answer position, including answer characters and terminal EOS.Source: §4.3 · p. 8

5 Related Work

The paper situates TANGO and WANGO among cross-token, parameter-generating, and linear-cost sequence models. Its central distinction is that source tokens generate feature-wise gates aggregated at destinations rather than general transformations or value vectors.

Cross-token gated computation

  • TANGO replaces separate self-attention and position-wise feed-forward processing with source-generated gates aggregated across visible tokens at each destination.This couples cross-token interaction with nonlinear gating in one update.Source: §5 · p. 9
  • Unlike conventional feed-forward gating, TANGO computes the gate from source representations and applies a weighted average of visible source gates at the destination.The destination’s projected features are rescaled feature-wise by the aggregated gate.Source: §5 · p. 9

Parameter sharing across depth

  • Parameter-sharing work studies repeated computation with shared blocks, whereas TANGO, WANGO, and Recurrent Transformer++ reuse one block across four applications in these experiments.Related settings include Universal Transformers, ALBERT, deep-equilibrium models, and looped language models.Source: §5 · p. 9

Sequence models with linear cost

  • Sparse attention, kernelized attention, linear attention, and recurrent state-space models reduce context cost through restricted interactions, accumulated statistics, or fixed-size states.GAU and FLASH provide gated-attention baselines with different sequence-length complexity.Source: §5 · p. 9
  • WANGO combines recent local pairwise interactions with positive-feature-map statistics summarizing older sources, yielding linear sequence-length computation for fixed window and feature dimensions.This places WANGO near sliding-window and linear-attention hybrids such as BASED.Source: §5 · p. 9

Conditional computation

  • Hypernetworks, Lambda layers, and fast-weight systems generate or store transformations, but TANGO generates SwiGLU gates rather than a general weight matrix.Mixture-of-experts models instead select among discrete or sparse learned modules.Source: §5 · p. 9

Formal proofs and generated mathematics

  • Lean and DeepMind Mathematics evaluate formal or generated mathematics, but this paper uses teacher-forced validation NLL rather than proof acceptance or generated-answer exact match.The reported tasks include Mathlib source-code modeling, proof completion, and mathematical answers.Source: §5 · p. 9

6 Discussion

TANGO achieves the lowest validation NLL across the tested benchmarks, while WANGO offers a linear-complexity alternative with favorable NLL at similar analytical cost to Recurrent Transformer++. These results require caution because operation counts are analytical, comparisons vary multiple architectural components, and the experiments do not establish compute-matched gains or scaling laws.

Validation NLL and computational cost

  • TANGO has the lowest validation NLL on FineWeb-Edu, Lean, and DeepMind Mathematics, but also the highest analytical operation count in the four-block experiments.For 8,192-token FineWeb-Edu sequences, TANGO requires 9.44T MACs versus 1.90T for WANGO and 1.94T for Recurrent Transformer++.Source: §6.1 · p. 10
  • WANGO has lower FineWeb-Edu validation NLL than Recurrent Transformer++ at a similar analytical operation count and lower NLL than FLASH.WANGO achieves this using local pairwise interactions and feature-map prefix sums for older source gates.Source: §6.1 · p. 10
  • Analytical operation counts compare arithmetic specified by model equations rather than wall-clock speed, GPU memory, or complete implementation costs.The counts omit normalization, masking, feature-map construction, and elementwise operations.Source: §6.1 · p. 10

Relation to self-attention and comparison architectures

  • TANGO and WANGO average source-produced SwiGLU gates, which rescale destination projected features instead of averaging self-attention value vectors.The resulting operator provides feature-wise scaling but does not generate arbitrary full projection matrices.Source: §6.2 · p. 10
  • The architecture comparisons vary several components simultaneously, so the NLL differences do not isolate cross-token gate aggregation as the cause.The comparison includes differences in attention, gating, parameter sharing, and block design.Source: §6.2 · p. 10

Limitations

  • The experiments match parameter counts and training conditions but do not match forward-pass operation counts or establish scaling laws.TANGO’s lower validation NLL should not be interpreted as a speed, memory, or compute-matched advantage.Source: §6.3 · p. 10
  • The study reports validation NLL and analytical operation counts, not generated-answer accuracy, proof acceptance, empirical training speed, or GPU memory.Its conclusions are limited to the tested model size, data splits, and training budgets.Source: §6.3 · p. 10

7 Conclusion

The paper replaces the standard self-attention and position-wise feed-forward pair with a cross-token gated update. TANGO has the best validation NLL across the tested benchmarks, while WANGO provides the strongest NLL among linear-complexity architectures without measuring downstream accuracy or hardware efficiency.

Conclusion

  • TANGO and WANGO replace self-attention plus the position-wise feed-forward network with one cross-token gated residual update.Each destination applies a weighted average of source SwiGLU gates to its own projected features.Source: §7 · p. 11
  • TANGO aggregates every causally visible source with quadratic sequence-length complexity, whereas WANGO uses a recent window and prefix statistics for linear complexity.WANGO is linear in sequence length when window and feature dimensions are fixed.Source: §7 · p. 11
  • At the evaluated parameter scale, TANGO has the lowest validation NLL on FineWeb-Edu, Lean, and DeepMind Mathematics.WANGO has the lowest NLL among tested linear-complexity architectures and beats Recurrent Transformer++ at a similar analytical operation count on FineWeb-Edu.Source: §7 · p. 11

Related papers

Loading 2608.22117v1…