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.

  • 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.

  • Results

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

  • 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.

  • Takeaways & Limitations

    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.

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.
  • 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.
  • 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.
  • 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.
  • 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.

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-conditioned gating: 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.
  • 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).
  • 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.
  • WANGO aggregation: 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.
  • 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.

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.
  • 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.
  • 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.
  • 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.
  • Benchmarks: 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.
  • Benchmarks: 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.
  • 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.

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.
  • 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.
  • Eight-block results: 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.
  • 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.
  • 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.
  • Lean and DeepMind Mathematics: 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.

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.
  • Cross-token gated computation: 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.
  • 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.
  • 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.
  • Sequence models with linear cost: 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.
  • 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.
  • 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.

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++.
  • Validation NLL and computational cost: 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.
  • Validation NLL and computational cost: 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.
  • 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.
  • Relation to self-attention and comparison architectures: 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.
  • 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.
  • Limitations: 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.

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.
  • Conclusion: 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.
  • Conclusion: 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.
Loading 2608.22117v1…