TIDE: Cada Capa Conoce el Token Subyacente al Contexto

TIDE: Every Layer Knows the Token Beneath the Context

May 7, 2026

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

cs.AI

Resumen

Revisitamos una elección de diseño universalmente aceptada pero poco examinada en todos los LLM modernos: un índice de token se busca una sola vez en la capa de *embedding* de entrada y luego se descarta permanentemente. Este supuesto de inyección única induce dos fallos estructurales: (i) el Problema del Token Raro, donde una distribución de vocabulario de tipo Zipf provoca que los *embeddings* de tokens raros estén crónicamente subentrenados al recibir una fracción de la señal de gradiente acumulada en comparación con los tokens comunes; y (ii) el Problema del Colapso Contextual, donde modelos con parámetros limitados mapean tokens distribucionalmente similares a estados ocultos indistinguibles. Como un intento de abordar ambos, proponemos TIDE, que aumenta el *transformer* estándar con una Memoria de *Embeddings*: un conjunto de K Bloques de Memoria independientes que mapean índices de token a vectores semánticos independientes del contexto, calculados una vez e inyectados en cada capa a través de un enrutador *softmax* condicionado por la profundidad con un banco nulo entrenable. Establecemos teórica y empíricamente los beneficios de TIDE para abordar los problemas asociados con la inyección única de identidad del token, así como para mejorar el rendimiento en múltiples tareas de modelado de lenguaje y tareas posteriores.

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.