TIDE: Jede Ebene kennt den Token unterhalb des Kontexts

TIDE: Every Layer Knows the Token Beneath the Context

May 7, 2026

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

cs.AI

Zusammenfassung

Wir hinterfragen eine allgemein akzeptierte, aber kaum untersuchte Designentscheidung in jedem modernen LLM: Ein Token-Index wird einmal in der Eingabe-Einbettungsschicht nachgeschlagen und dann dauerhaft verworfen. Diese Annahme einer einmaligen Injektion führt zu zwei strukturellen Problemen: (i) das Problem seltener Tokens, bei dem eine Zipf-verteilte Vokabularhäufigkeit dazu führt, dass Einbettungen seltener Tokens chronisch untertrainiert sind, da sie nur einen Bruchteil des kumulativen Gradientensignals im Vergleich zu häufigen Tokens erhalten; und (ii) das Problem des kontextuellen Kollapses, bei dem Modelle mit begrenzten Parametern distributionell ähnliche Tokens auf ununterscheidbare versteckte Zustände abbilden. Als Lösungsansatz für beide Probleme schlagen wir TIDE vor, das den Standard-Transformer um ein EmbeddingMemory erweitert: ein Ensemble von K unabhängigen MemoryBlocks, die Token-Indizes auf kontextfreie semantische Vektoren abbilden. Diese werden einmal berechnet und durch einen tiefenabhängigen Softmax-Router mit einer lernbaren Null-Instanz in jede Schicht injiziert. Wir belegen theoretisch und empirisch die Vorteile von TIDE bei der Lösung der Probleme, die mit der einmaligen Token-Identitätsinjektion verbunden sind, sowie die Leistungsverbesserung in mehreren Sprachmodellierungs- und Downstream-Aufgaben.

English

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.