TIDE : Chaque couche connaît le jeton sous-jacent au contexte

TIDE: Every Layer Knows the Token Beneath the Context

May 7, 2026

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

cs.AI

Résumé

Nous revisitons un choix de conception universellement accepté mais peu examiné dans chaque LLM moderne : un indice de token est recherché une seule fois au niveau de la couche d'embedding d'entrée puis définitivement abandonné. Cette hypothèse d'injection unique induit deux défaillances structurelles : (i) le Problème des Tokens Rares, où une distribution de type Zipf du vocabulaire entraîne un sous-entraînement chronique des embeddings des tokens rares, ceux-ci ne recevant qu'une fraction du signal de gradient cumulé par rapport aux tokens courants ; et (ii) le Problème de l'Effondrement Contextuel, où les modèles aux paramètres limités mappent les tokens distributionnellement similaires vers des états cachés indiscernables. Pour tenter de résoudre ces deux problèmes, nous proposons TIDE, qui augmente le transformeur standard avec une EmbeddingMemory : un ensemble de K MemoryBlocks indépendants qui mappent les indices de tokens vers des vecteurs sémantiques non contextuels, calculés une fois et injectés à chaque couche via un routeur softmax conditionné par la profondeur avec une banque nulle apprenable. Nous établissons théoriquement et empiriquement les bénéfices de TIDE pour résoudre les problèmes associés à l'injection unique de l'identité du token, ainsi que pour améliorer les performances sur de multiples tâches de modélisation du langage et tâches en aval.

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.