PaperlayerView exact revision on arXiv

Source-linked AI summary

TIDE: Every Layer Knows the Token Beneath the Context

Ajay Jaiswal, Lauren Hannah, Han-Byul Kim, Duc Hoang, Mehrdad Farajtabar, Minsik Cho

arXiv:2605.06216v1cs.CLcs.AIcs.LG

TL;DR

Modern transformers discard token identity after the input embedding, creating rare-token under-training and contextual collapse. TIDE adds persistent token-indexed semantic memory to every layer, and the paper reports consistent performance improvements across language-modeling and downstream evaluations. The architecture is also theoretically shown to amplify per-token gradient signal by K and address the identified failure modes.

  • Problem

    A single input-layer token-identity injection leaves rare-token embeddings under-trained and can make semantically distinct tokens indistinguishable in deeper layers.

  • Method

    TIDE adds EmbeddingMemory, an ensemble of K token-indexed MemoryBlocks whose context-free semantic vectors are routed into every transformer layer.

  • Results

    TIDE consistently improves performance across language-modeling datasets and downstream tasks, while theory shows a K-fold amplification of per-token cumulative gradient signal.

  • Takeaways & Limitations

    Persistent token-specific signals provide an architectural way to address the issues associated with single-token identity injection.

  • Takeaways & Limitations

    TIDE's contextual-collapse guarantee is bounded by fixed upstream representations and the practical constraint on the FFN Lipschitz constant.

Abstract

from arXiv · show

We revisit a universally accepted but under-examined design choice in every modern LLM: a token index is looked up once at the input embedding layer and then permanently discarded. This single-injection assumption induces two structural failures: (i) the Rare Token Problem, where a Zipf-type distribution of vocabulary causes rare-token embeddings are chronically under-trained due to receiving a fraction of the cumulative gradient signal compared to common tokens; and (ii) the Contextual Collapse Problem, where limited parameters models map distributionally similar tokens to indistinguishable hidden states. As an attempt to address both, we propose TIDE, which augments the standard transformer with EmbeddingMemory: an ensemble of K independent MemoryBlocks that map token indices to context-free semantic vectors, computed once and injected into every layer through a depth-conditioned softmax router with a learnable null bank. We theoretically and empirically establish the benefits of TIDE in addressing the issues associated with single-token identity injection as well as improve performance across multiple language modeling and downstream tasks.

1 Introduction

Modern transformers inject token identity only once, leaving rare-token embeddings under-trained and allowing semantically distinct tokens to collapse into similar hidden states. TIDE addresses these failures with persistent token-indexed semantic memory injected throughout the network, with reported gains across scales, language-modeling datasets, and downstream tasks.

  • Motivation: Token identity is looked up at the input embedding layer and then discarded, so later layers operate only on contextualized hidden states.This single-injection assumption underlies both failure modes described in the section.
  • Rare Token Problem: Rare-token embeddings remain under-trained because their cumulative SGD gradient signal is proportional to frequency under Zipf-like vocabularies.The paper identifies rare named entities, technical terms, and low-frequency morphological forms as examples.
  • Contextual Collapse Problem: Semantically distinct tokens in similar syntactic environments can produce indistinguishable hidden states when intermediate layers cannot re-consult token identity.The paper attributes this contextual collapse to representational overloading in FFNs and contextual mixing in the residual stream and attention outputs.
  • Empirical Evidence: Figure 1 shows rare-token embedding norms are noisier and decline during continued training while common-token norms increase.The figure reports mean norms, norm distributions, and bin-wise growth rates across intermediate checkpoints.
  • TIDE Architecture: TIDE adds EmbeddingMemory, an ensemble of K independent MemoryBlocks that map token indices to context-free semantic vectors for injection into every transformer layer.The memory is computed once per forward pass and routed through a per-layer softmax mechanism conditioned on the post-attention hidden state.
  • Contributions: TIDE amplifies per-token cumulative gradient signal by K and consistently improves performance across 350M-to-1B models, language-modeling datasets, and downstream tasks.The reported evaluations include Wikitext, PubMed, DCLM, HellaSwag, ARC, and PIQA.

2 When Context is Not Enough: Diagnosing Standard Transformers

Standard transformers suffer from gradient starvation for rare tokens and contextual collapse when distinct tokens appear in nearly identical contexts. These failures limit token-specific representation and cannot always be repaired by downstream FFNs.

  • 2.1 The Rare Token Problem.: Rare-token embeddings receive far fewer gradient updates than common-token embeddings under minibatch SGD.The expected update count scales with token frequency, and the Wikitext-103 instantiation reports a six-orders-of-magnitude gradient-signal disparity.
  • 2.1 The Rare Token Problem.: Over 200B training tokens, rare-token norms decline while common-token norms continuously increase, showing that the disparity grows during training.The reported trend is not limited to a cold-start effect.
  • 2.2 Contextual Collapse and the FFN’s Blind Spot.: Contextual collapse occurs when distinct tokens in nearly identical syntactic environments produce similar hidden states across transformer layers.The examples include grammatical homophones, numeric identity tokens, and rare domain-specific synonyms.
  • 2.2 Contextual Collapse and the FFN’s Blind Spot.: Figure 2 shows near-zero hidden-state distances for most layers, with numerical tokens exhibiting especially severe collapse even near the final layer.The heatmap summarizes token-pair distances across template sentences from three collapse categories.
  • 2.2 Contextual Collapse and the FFN’s Blind Spot.: When collapsed-token separation exceeds the FFN’s Lipschitz capacity, no FFN can approximate the desired token-specific mapping arbitrarily well, regardless of width.The bound attributes the limitation to fixed upstream representations and the continuous FFN transformation.
  • 2.2 Contextual Collapse and the FFN’s Blind Spot.: TIDE addresses the missing identity signal by adding an EmbeddingMemory module whose K MemoryBlocks inject token-indexed signals into every layer.A per-layer router combines the memory-block outputs and adds them to the residual stream.

3 TIDE: Token Identity Delivered Everywhere

TIDE preserves token identity throughout the transformer by adding a dedicated, token-indexed semantic memory pathway alongside the contextual residual stream. Its K independent MemoryBlocks support repeated layer-wise identity conditioning, amplify rare-token gradient signals, and improve token separation while retaining baseline expressivity.

  • Architecture: TIDE counters single-injection design by making token identity directly accessible at every transformer depth.The memory signal remains independent of contextual mixing in the residual stream.
  • Architecture: Each MemoryBlock independently maps token indices to context-free vectors, with K blocks stacked into one memory tensor shared across all layers.The blocks use separate embedding tables without parameter sharing.
  • Architecture: A depth-conditioned softmax router combines memory blocks and a learnable null bank, then adds the selected memory vector alongside the FFN pathway.The null bank provides an effective off switch, while the memory contribution remains token-indexed rather than hidden-state-indexed.
  • Overhead: TIDE’s per-layer overhead is limited to a softmax router and weighted sum, while static memory tables can be 4-bit quantized and offloaded to SSD.The reported VRAM footprint remains similar to LLaMA-Base-1B, at 1.03 GB in 8-bit, while SSD usage grows to 3.152 GB at K=24.
  • Expressivity: TIDE can approximate the standard transformer function class to arbitrary precision through its learnable null bank.Assigning weight to the null bank suppresses active memory contributions.
  • Gradient amplification: K independent pathways provide K-fold amplification of cumulative gradient signal for each token relative to the standard single-embedding baseline.When a token appears, its gradient flows through all K independent embedding tables simultaneously.
  • Rare-token benefits: 0.704 nats (9.0% relative) is TIDE’s rarest-decile loss reduction, compared with 0.068 nats (2.4%) for the most frequent decile.TIDE outperforms LLaMa-Base-1B on every frequency decile, with gains decreasing from rare to common tokens.
  • Contextual separation: TIDE increases layer-wise hidden-state separation across all three contextual-collapse categories, especially from middle to terminal layers.Numerical tokens are identified as the predominant beneficiaries of token-identity injection.

4 Experiments and Ablation Studies

Experiments show that TIDE improves language modeling and downstream performance, with especially large gains for rare tokens. MemoryBlocks learn complementary token-identity representations, while routing varies by token frequency.

  • Perplexity and Training Dynamics: TIDE variants strictly outperform LLaMA-Base-1B across Wikitext, PubMed, and DCLM as K increases from 2 to 24.With 100B training tokens, TIDE using 2–4 MemoryBlocks matches the baseline perplexity reached with 200B tokens.
  • Influence of K across Rare, Mid, and Common Tokens: 0.421 nats: rare-token loss falls from 6.671 to 6.250 as K increases from 0 to 24, versus a 0.075-nat common-token reduction.The rare-token absolute gain is 5.6× larger, and the per-block marginal benefit is 3.7× steeper than for common tokens.
  • Influence of K across Rare, Mid, and Common Tokens: ∼55% of the total rare-token improvement at K = 24 is achieved with only K = 2 MemoryBlocks.The passage suggests that most of the rare-token benefit can be obtained with modest 2–4-block configurations.
  • TIDE and Downstream Task Performance: +2.3% absolute gain: at 1B parameters, TIDE raises average zero-shot accuracy from 61.4 to 63.7 across eight benchmarks.Performance improves monotonically with K on perplexity and six of eight downstream tasks.
  • Distance between Primary Embedding and MemoryBlocks: MemoryBlocks encode complementary token-identity information rather than replicating the primary embedding subspace.Mean cosine distance from the primary embedding ranges from 0.65 to 0.99, while inter-block distance is smaller but non-collapsed.
  • Bin-wise Router Statistics for MemoryBlocks and the NULL Bank: The NULL-bank weight rises from 0.530 for the rarest decile to 0.889 for the most common, allocating more memory mass to rare tokens.Routing is also non-uniform across blocks: distinct banks specialize in different frequency regimes.

5 Conclusion

The conclusion presents TIDE as an architectural response to rare-token gradient starvation and contextual collapse. It combines persistent token-indexed memory with layer-wise routing and reports consistent gains across language modeling and downstream tasks.

  • Conclusion: TIDE addresses gradient starvation of rare tokens and contextual collapse of semantically distinct tokens in standard transformers.These are described as empirically established failure modes associated with single-token identity injection.
  • Conclusion: EmbeddingMemory uses K independent MemoryBlocks to map token indices to semantic vectors injected at every layer through a depth-conditioned router with a NULL bank.The memory signal is persistent, token-specific, and complementary to the contextual residual stream.
  • Conclusion: TIDE consistently improves performance across multiple language-modeling and downstream tasks at different model scales.The conclusion summarizes the paper’s theoretical and empirical evaluation without restricting the claim to a single dataset or benchmark.
  • Motivation: Zipfian long-tail token frequencies leave rare embeddings under-trained because their cumulative gradient signal is proportional to occurrence frequency.Subword tokenization and embedding sharing are described as insufficient to resolve this fundamental long-tail issue.

C Gradient Starvation Bound: Derivation of Equation (2.1)

The derivation bounds how often a token embedding receives nonzero gradient and shows that this frequency dependence produces a severe rare–common token disparity. A WikiText-103 instantiation quantifies the gap under explicit assumptions.

  • Probability of Token Appearance: Pr[v ∈ batch_s] = 1 − (1 − f_v)^BT ≤ f_v · B · T bounds the probability that token v appears in a batch.The bound follows from the Bernoulli inequality for f_v ∈ [0, 1].
  • Cumulative Gradient Bound: τ · f_v · B · T · G^2 bounds the cumulative squared gradient for token v over τ training steps.The bound uses the per-step gradient bound and the indicator of whether v appears in each batch.
  • Rare–Common Ratio: O(ε/c) bounds the ratio of rare-token to common-token cumulative gradient under rare frequency ε and common frequency lower bound c.The derivation treats BT, gradient bounds, and the common-token appearance constant as fixed positive quantities.
  • Concrete Evaluation: κ ≈ 1 for the common-token appearance probability in the WikiText-103 instantiation with c = 8.3 × 10−3, B = 8, and T = 2048.The calculation gives 1 − e^−136 ≈ 1.
  • Concrete Evaluation: 10−5 times: under G^2/G^2_min = 10, a Bin-0 hapax embedding’s gradient signal is bounded above by this fraction of a Bin-9 common token’s signal.The result corresponds to a five-orders-of-magnitude disparity over the same training run.

D Full Proof of Proposition 2.2

The proof formalizes contextual collapse as a limitation of Lipschitz FFNs: nearly identical upstream hidden states constrain their outputs even when the task requires separated token representations.

  • Approximation Bound: When target separation C exceeds L_FFNδ, at least one token’s approximation error must remain positive.The proof applies the triangle inequality and the fact that the maximum of two nonnegative errors is at least half their sum.
  • Lipschitz Constraint: L_FFNδ bounds the distance between FFN outputs for a collapsed token pair, regardless of the chosen FFN weights.The bound follows from the FFN’s Lipschitz property.
  • Setup: When attention produces hidden-state distance δ between tokens u and v, the FFN receives those fixed inputs and cannot choose different ones.This establishes input proximity as an upstream constraint on the FFN’s ability to separate tokens.
  • Practical Limitation: Increasing L_FFN to overcome collapse can amplify every input perturbation, causing exploding gradients, training instability, and degraded performance on non-collapsed tokens.The practical change limit is therefore constrained by effects beyond the collapsed token pair.

E Full Proof of Proposition 3.1

The proof shows that a sufficiently large null-bank logit can suppress the active memory-bank contribution arbitrarily closely, while the required threshold remains finite.

  • Null-bank suppression: The memory norm is bounded by KC/(K + e^s), where C is the maximum MemoryBlock output norm.This follows from the softmax weights and the uniform bound C on MemoryBlock outputs.
  • Null-bank suppression: The total active-bank weight vanishes as the null logit s approaches infinity.The bound K/(K + e^s) tends to zero.
  • Uniform approximation: For any ϵ ∈ (0, C), solving the bound yields a finite threshold s∗ beyond which the memory contribution is at most ϵ uniformly over tokens and layers.The proof explicitly states that s∗ is finite for every admissible ϵ.
  • Uniform approximation: A single large null logit suppresses all K active banks jointly, reducing the suppression degree of freedom to one scalar.The null bank therefore provides coordinated control over the entire active memory ensemble.

F Full Proof of Proposition 3.2: K-Pathway Gradient Amplification

The proof establishes K-pathway gradient amplification: each independent MemoryBlock receives layer-wise gradient contributions whenever a token appears, and summing across blocks yields a K-fold lower bound under the stated assumptions.

  • Proof setup: The simplified Proposition 3.2 router excludes the null bank and considers K active MemoryBlocks.This is the proof setting used for the gradient-amplification result.
  • Per-block gradients: A MemoryBlock embedding receives zero gradient when its token is absent from the batch and layer-wise upstream gradients when the token is present.Because each block enters every layer, its gradient accumulates across all L transformer layers.
  • Per-block gradients: Under the non-degeneracy assumption, each independent block satisfies the same positive gradient lower-bound structure for tokens observed in training.The assumption requires at least one layer with a nonzero relevant contribution whenever the token remains imperfectly learned.
  • Amplification result: Summing the independent block contributions gives a K-fold amplification relative to the analogous single-block lower bound.The result is stated under the proof’s G2/G2 assumptions.

G Additional Details for Proposition 3.3: Contextual Collapse and TIDE’s K-MemoryBlocks

TIDE’s token-indexed MemoryBlocks can separate collapsed token pairs independently of their similar contextual hidden states, unlike the FFN’s Lipschitz-bounded response. Layer ablations further show that memory contributions are intermittent, with especially strong early and secondary effects.

  • Contextual collapse: MemoryBlock outputs depend on discrete token identity rather than the hidden-state gap δ between the pair.The token index is absent from the contextual FFN computation but directly selects each MemoryBlock row.
  • Contextual collapse: The FFN output separation is bounded by L_FFNδ, limiting its ability to resolve collapsed pairs without affecting other inputs.Increasing L_FFN amplifies all input perturbations and can destabilize non-collapsed tokens.
  • Layer-wise contribution: Dropping layer 0 raises perplexity by more than 103% on every dataset and by 1.09 × 10^6% on PubMed.Layer 1 also remains load-bearing, with degradation of +8.1% to +12.9% across datasets.
  • Layer-wise contribution: Memory contribution is intermittent rather than monotone, with a secondary degradation peak at layer 13.The authors interpret this as early identity information persisting before requiring intermittent refresh.
  • Storage considerations: TIDE-8E-1B requires approximately 4.2 GB of F16 storage for EmbeddingMemory, motivating quantization and low-rank compression investigations.The storage study examines both quantization and low-rank compression as ways to address this cost.

J.2 Low Rank Compression of Token-wise MemoryBlocks.

Uniform low-rank compression preserves perplexity relatively well at modest reductions but causes sharp degradation beyond 70%; MemoryBlocks also provide especially complementary neighbor information for rare tokens.

  • Rank-storage trade-off: For TIDE-8E-1B, rank r ≤ 2015 is sufficient for any uniform reduction of at least approximately 2% to reduce storage below dense parameterization.The bound uses vocabulary size 128,256 and hidden and bottleneck dimensions of 2048.
  • Perplexity under compression: Perplexity degradation remains almost flat at 10–30% reduction, then grows gradually at 40–60% reduction.At 10–30%, each MemoryBlock can shrink to 71% of its dense parameter count.
  • Perplexity under compression: Beyond 70% reduction, relative perplexity degradation reaches 587% on WikiText-2, 484% on DCLM, and 657% on PubMed.The curves bend sharply upward in this high-compression regime.
  • Semantic complementarity: Rare-token neighbor sets in MemoryBlocks are consistently more disjoint from the primary embedding’s neighbors than common-token sets.This indicates complementary, non-overlapping information for rare tokens rather than merely a global statistical artifact.

M Limitations and Future Work

TIDE’s evaluation has practical storage, scope, and interpretability boundaries. The authors identify compression needs, untested training regimes, and unresolved questions about MemoryBlock specialization.

  • Storage overhead: EmbeddingMemory storage overhead scales linearly with K, requiring compression strategies for deployments with strict storage budgets.The tables are static and quantization friendly, but their SSD footprint still grows with the number of MemoryBlocks.
  • Evaluation scope: TIDE’s benefits remain unexplored for longer training horizons, instruction tuning, and RLHF.Experiments cover 750M–3B parameter models trained on 200–500B DCLM tokens, with evaluations on WikiText, PubMed, DCLM, and eight zero-shot benchmarks.
  • Interpretability: The paper does not provide a principled account of what each MemoryBlock learns.Router statistics and nearest-neighbor analyses suggest specialization across frequency regimes, motivating finer-grained interpretability studies.
Loading 2605.06216v1…