Technical Summary: TANGO and WANGO

Problem Statement

Standard Transformer architectures separate cross-token interaction (handled by self-attention) from position-wise nonlinear transformation (handled by feed-forward networks). This separation requires two distinct sublayers per block. While effective, this design may be suboptimal for efficiency and parameter utilization. Furthermore, existing linear-complexity sequence models often rely on approximations that may sacrifice the richness of cross-token interaction or require complex gating mechanisms that do not fully integrate the nonlinear transformation with the attention mechanism.

Methodology

Core Architecture: Token-Aggregated Nonlinear Gating Operators (TANGO)

The paper introduces TANGO, a decoder block that replaces the standard self-attention and position-wise feed-forward sublayers with a single cross-token gated residual update.

1. Source-Conditioned SwiGLU Gating:

   • Unlike a conventional SwiGLU (Swish-Gated Linear Unit) where a gate and features are computed from the same token, TANGO allows a gate computed at a source position ($x_j$) to modulate projected features at a destination position ($x_i$).
   • Each position $i$ produces a gate vector $g_i = \text{SiLU}(W_g x_i)$ and a projected feature vector $v_i = W_v x_i$.
   • For a destination $i$, the model computes a content-dependent weighted average of the gate vectors from all causally visible sources ($j \le i$), denoted as $\bar{g}_i$.
   • The residual update is computed as:
     $$h^+_i = h_i + W_o(\bar{g}_i \odot v_i)$$
   • Here, the aggregated gate $\bar{g}_i$ acts as a diagonal scaling operator on the destination's projected features $v_i$ before the final output projection $W_o$.

2. Weight Aggregation Mechanisms:
   The paper proposes two variants differing in how they compute the weights for aggregating source gates:

   • TANGO (Full-Prefix): Computes softmax weights based on scaled cosine similarities between every destination and every causally visible source. This results in quadratic complexity ($O(T^2)$) with respect to sequence length $T$.
   • WANGO (Windowed Aggregation): A linear-complexity variant ($O(T)$ for fixed window/feature dimensions).
     • Recent Sources (Window): Uses the same exponential cosine-similarity scores as TANGO for sources within a trailing window $L_i$.
     • Older Sources (Prefix): Uses a deterministic positive feature map $\phi(u)$ to compute scores. This allows the model to maintain running sums of transformed keys and their outer products with gate vectors. The contribution of older sources is computed via these running sums, avoiding the need to re-evaluate all past pairs.

Experimental Setup

The author compares TANGO and WANGO against four baselines:

• Recurrent Transformer++: A standard Transformer block reused four times with shared parameters.
• Untied Transformer++: Four independently parameterized Transformer blocks.
• Full-attention GAU (Gated Attention Unit): Four independent blocks using full pairwise attention.
• FLASH: Four independent blocks combining local pairwise attention with a global linear-attention summary.

Constraints:

• All models are matched for nonembedding parameter counts (approx. 44.3M distinct parameters).
• Models are trained with identical data orders, initialization seeds, and training budgets.
• Experiments are conducted on FineWeb-Edu (natural language), Lean (formal proofs and source code), and DeepMind Mathematics.
• Comparisons are made with both 4-block and 8-block sequential applications.

Key Contributions

1. TANGO Model: Introduces a unified block replacing self-attention and feed-forward networks with a single gated update where source tokens produce SwiGLU gates that modulate destination features.
2. WANGO Model: Develops a linear-complexity variant that retains exact pairwise interactions within a local window and uses feature-map prefix sums for older sources, maintaining the same parameterization as TANGO.
3. Rigorous Comparison: Provides a controlled comparison of six architectures under strictly matched parameter counts, training objectives, and data ordering, isolating the effects of architectural design from training variance.

Results

Performance on Benchmarks

• FineWeb-Edu:
  • TANGO achieved the lowest mean validation Negative Log-Likelihood (NLL) overall (3.293).
  • WANGO achieved the lowest NLL among all architectures with linear sequence-length scaling (3.355), outperforming both FLASH (3.465) and Recurrent Transformer++ (3.392).
  • Notably, Wango achieved lower NLL than Recurrent Transformer++ at a nearly identical analytical forward-pass multiply–accumulate (MAC) count (1.90T vs 1.94T for an 8,192-token sequence).
• Lean (Source Code & Proof Completion):
  • TANGO obtained the lowest joint NLL (1.719) and the lowest NLL for both source-code modeling and proof completion separately.
  • WANGO (2.271) outperformed FLASH (2.377) but did not outperform Recurrent Transformer++ (2.530) in the 4-block experiment; however, in the 8-block experiment, WANGO (2.231) did outperform Recurrent Transformer++ (2.364).
• DeepMind Mathematics:
  • TANGO achieved the lowest validation NLL (0.473).
  • WANGO (0.487) performed better than Recurrent Transformer++ (0.537) and Untied Transformer++ (0.556).

Computational Complexity

• TANGO: Quadratic in sequence length ($O(T^2)$). For an 8,192-token sequence, it requires ~9.44T MACs.
• WANGO: Linear in sequence length ($O(T)$) for fixed dimensions. For the same sequence, it requires ~1.90T MACs.
• Recurrent Transformer++: Quadratic complexity but lower MACs than TANGO (~1.94T) due to parameter sharing and lack of full cross-token gating aggregation.

Significance and Claims

The paper claims that TANGO achieves state-of-the-art validation likelihood among the tested architectures across natural language, formal code, and mathematical reasoning tasks at the evaluated parameter scale.

The author emphasizes that WANGO demonstrates that linear-complexity models can outperform quadratic-complexity baselines (like Recurrent Transformer++) when matched for analytical operation counts and parameter size, specifically on FineWeb-Edu. Specifically, WANGO's ability to aggregate SwiGLU gates via running sums allows it to maintain high performance with linear scaling.

The author explicitly states the limitations of their claims:

• The results are based on validation NLL under teacher forcing and do not measure generated-answer accuracy, proof acceptance rates, or wall-clock training speed.
• The experiments do not establish scaling laws beyond the tested parameter count (approx. 70M total parameters).
• The lower NLL of TANGO should not be interpreted as a speed or memory advantage, as it incurs significantly higher computational costs than WANGO or Recurrent Transformer++.

In summary, the paper presents a novel architectural block that unifies attention and feed-forward mechanisms through cross-token gating, demonstrating that this approach yields superior predictive likelihood, with a linear-complexity variant (WANGO) offering a competitive alternative for long-context applications.