TIDE: Every Layer Knows the Token Beneath the Context

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.

📋 한 줄 요약

**[LLM/아키텍처]** 토큰 인덱스를 입력 임베딩에서 한 번만 주입하던 관행을 깨고, 모든 레이어에 토큰 정체성을 주입하는 EmbeddingMemory 구조 TIDE를 제안한다.

🎯 핵심 기여도

현대 LLM의 보편적 설계 결함인 "single-injection assumption"을 분석하여 Rare Token Problem과 Contextual Collapse Problem을 구조적으로 규명
K개의 독립적인 MemoryBlock 앙상블로 구성된 EmbeddingMemory 모듈 제안
depth-conditioned softmax router와 학습 가능한 null bank를 통해 레이어별로 토큰 의미를 재주입
다수의 언어 모델링 및 다운스트림 태스크에서 성능 향상을 이론·실증적으로 검증

💡 핵심 아이디어

토큰 인덱스가 입력 임베딩 단계에서 한 번만 조회된 후 영구히 폐기되는 기존 가정은 희귀 토큰의 그래디언트 부족(Zipf 분포)과 분포적으로 유사한 토큰의 hidden state 붕괴를 야기한다. TIDE는 토큰 정체성을 모든 레이어로 직접 다시 흘려보내 이 두 문제를 동시에 해결한다.

🔬 기술적 접근법

**모델/방법론**: 표준 Transformer에 EmbeddingMemory 모듈을 추가한 TIDE 아키텍처
**핵심 기법**: K개의 MemoryBlock 앙상블 + depth-conditioned softmax router + learnable null bank를 활용한 layer-wise context-free 의미 벡터 주입

📊 주요 결과

다양한 언어 모델링 벤치마크에서 베이스라인 대비 일관된 성능 향상
희귀 토큰 학습 부족과 contextual collapse 현상 모두에서 개선 확인
다운스트림 태스크에서도 성능 향상 입증

💭 의의 및 한계

**의의**: 모든 모던 LLM이 당연시하던 단일 임베딩 주입 가정을 재검토하여, 작은 모델에서도 표현력을 보강하는 일반적 해법을 제시. **한계**: K개의 MemoryBlock 추가에 따른 파라미터·연산 오버헤드와 대규모 모델에서의 스케일링 효과는 추가 검증이 필요.

🚀 실용적 활용

어휘 분포가 long-tail인 다국어/도메인 특화 LLM의 희귀 토큰 표현 강화
파라미터 제약이 있는 소형 LLM의 contextual capacity 보완
토큰 정체성 보존이 중요한 코드·기호·구조화 입력 모델링