TIDE: 모든 레이어가 문맥 아래의 토큰을 알고 있다

현대 LLM은 입력 단계에서만 토큰의 정체성을 확인하고 이후 레이어에서는 문맥 정보에만 의존하여, 자주 나오지 않는 희귀 토큰을 제대로 학습하지 못하거나 비슷한 문맥에서 토큰을 구분하지 못하는 한계가 있다. TIDE는 모든 레이어에 토큰 고유의 정보를 직접 주입하여 희귀 단어 처리 능력을 대폭 향상시키고 모델의 전반적인 성능을 개선한다.

기존 Transformer는 문장을 읽을 때 처음에만 각 단어가 무엇인지 확인하고, 이후 수십 개의 레이어를 거치는 동안에는 주변 단어들과 섞인 '문맥' 정보만을 처리한다. 이 과정에서 '사과'와 '배'가 비슷한 문장에서 쓰이면 깊은 레이어에서는 두 단어의 차이를 잊어버리는 '문맥적 붕괴'가 발생하며, 특히 드물게 등장하는 단어는 학습 신호(Gradient)를 충분히 받지 못해 제대로 학습되지 않는다.

TIDE는 이 문제를 해결하기 위해 모든 레이어에 '단어 신분증'을 다시 보여주는 방식을 채택했다. 각 레이어마다 현재 처리 중인 단어가 원래 무엇이었는지 알려주는 별도의 메모리 경로를 만들고, 이를 기존 문맥 정보와 합쳐서 계산한다. 이는 마치 시험을 볼 때 문제의 앞부분만 보고 푸는 것이 아니라, 풀이 과정 내내 문제의 핵심 키워드를 옆에 적어두고 참고하는 것과 같다.

결과적으로 모델은 문맥이 비슷하더라도 단어 고유의 의미를 끝까지 유지할 수 있게 된다. 실험 결과, 가장 드물게 등장하는 단어들에 대한 예측 정확도가 크게 향상되었으며, 전체적인 언어 모델링 성능도 기존 방식보다 적은 학습량으로 더 높은 수준에 도달했다.

TIDE는 표준 Transformer 블록에 EmbeddingMemory라는 병렬 경로를 추가한다. 이 모듈은 K개의 독립적인 MemoryBlock으로 구성되며, 각 블록은 토큰 인덱스를 입력받아 문맥과 무관한 고유의 시맨틱 벡터를 생성한다. [토큰 인덱스 v → Embedding Lookup → RMSNorm → Mk(v)] 과정을 통해 각 블록의 출력이 계산된다.

각 레이어에서는 현재의 숨겨진 상태를 바탕으로 어떤 MemoryBlock을 얼마나 사용할지 결정하는 라우팅 메커니즘을 사용한다. [Post-Attention 상태 n → Linear Router → Softmax → 라우팅 가중치 α] 연산을 수행한다. Softmax를 통해 출력된 가중치들의 합은 항상 1이 되며, 이는 각 메모리 블록의 중요도를 확률적으로 배분하는 의미를 갖는다.

최종적으로 선택된 메모리 값들은 기존 FFN의 출력과 더해져 잔차 연결(Residual Stream)에 주입된다. [Σ(α * Mk) → m(v) → h + FFN + m(v)] 순서로 계산이 이루어지며, 이를 통해 토큰 고유의 정체성 신호가 문맥 정보와 독립적으로 유지되면서도 상호 보완적으로 작용하게 된다.

LLaMa-Base 1B 모델과 비교했을 때, TIDE는 Wikitext, PubMed, DCLM 등 주요 데이터셋에서 일관된 성능 향상을 보였다. 특히 1B 규모에서 평균 벤치마크 점수가 61.4%에서 63.7%로 상승했다. 희귀 토큰(Rare Tokens)에 대한 Cross-Entropy 손실값은 기존 모델 대비 최대 9.0% 감소하여, 빈도가 낮은 단어 처리에 탁월한 효과가 있음을 입증했다.

효율성 측면에서 TIDE는 100B 토큰 학습만으로도 기존 Transformer가 200B 토큰을 학습했을 때의 성능을 추월하는 빠른 수렴 속도를 보였다. 또한 MemoryBlock의 개수 K가 증가함에 따라 성능이 단조 증가하는 경향을 확인했으며, K=24 설정에서도 VRAM 사용량은 8비트 기준 1.03GB 수준으로 유지하면서 추가 파라미터는 SSD에 저장하여 효율적으로 운영 가능함을 보여주었다.

TIDE 아키텍처는 'Single-injection assumption'이라는 기존 Transformer의 설계적 한계를 지적한다. 기존 방식은 입력층에서만 토큰 임베딩을 참조하고 이후에는 폐기하지만, TIDE는 이를 모든 레이어에서 재참조할 수 있도록 구조화했다. 이는 수학적으로 FFN의 Lipschitz 연속성 제약으로 인해 발생하는 정보 손실을 우회하는 효과를 준다.

EmbeddingMemory는 K개의 독립적인 경로를 제공하여 그래디언트 기아(Gradient Starvation) 문제를 해결한다. 이론적으로 TIDE는 표준 Transformer 대비 누적 그래디언트 신호를 K배 증폭시킬 수 있음을 증명했다. 또한 'Null Bank' 개념을 도입하여, 모델이 특정 레이어에서 토큰 정체성 정보가 필요 없다고 판단할 경우 메모리 주입을 스스로 차단할 수 있는 학습 가능한 스위치를 제공한다.

추론 시에는 메모리 테이블을 4비트로 양자화하여 SSD에 저장하고, 필요할 때 비동기적으로 프리페칭(Asynchronous Prefetching)하는 방식을 제안하여 하드웨어 요구 사항을 낮췄다. 이는 대규모 파라미터를 효율적으로 관리하면서도 레이어별로 최적화된 토큰 정보를 공급할 수 있는 실용적인 아키텍처이다.

EmbeddingMemory 테이블이 정적이고 양자화에 유리함에도 불구하고, SSD 점유 공간이 K 값에 비례하여 선형적으로 증가하는 저장 공간 오버헤드가 존재한다. 또한 본 연구는 3B 파라미터 규모까지만 실험되었으며, 더 큰 규모의 모델이나 지시어 튜닝(Instruction Tuning), RLHF 이후의 성능 변화에 대해서는 추가 검증이 필요하다.