TIDE: Cada Camada Conhece o Token Sob o Contexto

Resumo

Revisitamos uma escolha de projeto universalmente aceita, mas pouco examinada, em todos os LLMs modernos: um índice de token é consultado uma única vez na camada de incorporação de entrada e depois descartado permanentemente. Essa premissa de injeção única induz a duas falhas estruturais: (i) o *Problema do Token Raro*, onde uma distribuição Zipfiana do vocabulário faz com que as incorporações de tokens raros sejam cronicamente subtreinadas, por receberem uma fração do sinal gradiente cumulativo em comparação com tokens comuns; e (ii) o *Problema do Colapso Contextual*, onde modelos com parâmetros limitados mapeiam tokens distributionalmente semelhantes para estados ocultos indistinguíveis. Como uma tentativa de abordar ambos, propomos o TIDE, que amplia o *transformer* padrão com uma *EmbeddingMemory*: um conjunto de K *MemoryBlocks* independentes que mapeiam índices de token para vetores semânticos livres de contexto, computados uma vez e injetados em todas as camadas por meio de um roteador *softmax* condicionado à profundidade com um banco nulo aprendível. Estabelecemos teórica e empiricamente os benefícios do TIDE ao abordar os problemas associados à injeção única da identidade do token, bem como ao melhorar o desempenho em múltiplas tarefas de modelagem de linguagem e tarefas *downstream*.

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.