TIDE: Каждый слой знает токен под контекстом

TIDE: Every Layer Knows the Token Beneath the Context

May 7, 2026

Авторы: Ajay Jaiswal, Lauren Hannah, Han-Byul Kim, Duc Hoang, Mehrdad Farajtabar, Minsik Cho

cs.AI

Аннотация

Мы пересматриваем повсеместно принятый, но недостаточно изученный конструктивный выбор в каждой современной крупной языковой модели (LLM): индекс токена извлекается единожды на уровне входного эмбеддинга и затем окончательно отбрасывается. Это предположение о единичной инъекции порождает два структурных недостатка: (i) Проблему редких токенов, когда Zipf-подобное распределение словаря приводит к хроническому недообучению эмбеддингов редких токенов из-за получения лишь доли совокупного градиентного сигнала по сравнению с частотными токенами; и (ii) Проблему контекстуального коллапса, когда модели с ограниченным числом параметров отображают распределительно схожие токены в неразличимые скрытые состояния. В попытке решить обе проблемы мы предлагаем TIDE, который расширяет стандартный трансформер с помощью EmbeddingMemory: ансамбля из K независимых блоков памяти (MemoryBlocks), которые отображают индексы токенов в контекстно-независимые семантические векторы, вычисляемые один раз и инжектируемые в каждый слой через условный на глубине софтмакс-маршрутизатор с обучаемым нулевым банком. Мы теоретически и эмпирически доказываем преимущества TIDE в решении проблем, связанных с единичной инъекцией идентичности токена, а также в повышении производительности в различных задачах языкового моделирования и последующих прикладных задачах.

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.