Mela: Transformation Hypothesis에 기반한 Test-Time Memory Consolidation

Transformer 기반 언어 모델의 긴 시퀀스 처리 한계를 다층 기억 체계로 해소한다. Transformation Hypothesis와 cross-frequency coupling에서 영감을 얻은 Hierarchical Memory Module(HMM)과 Hierarchical Latent Recursion(HLR)를 통해, 저주기 메모리(L-module)와 고주기 메모리(H-module)가 상호작용하며 컨솔리데이션을 수행한다. MemStack은 중간 컨솔리데이션 상태를 디코더의 초반 층에 주입해 긴 컨텍스트에서도 안정적인 예측을 가능하게 한다.

단계 1: Transformer의 Self-attention은 시퀀스 내 모든 토큰 쌍의 관계를 계산하므로 길이가 늘어나면 연산량이 증가한다. 단계 2: transformation hypothesis에 따라 메모리는 고정된 기록이 아니라 상황에 맞게 변화하는 표현으로 재구성되며, 서로 다른 시간 규모의 시스템이 협력해야 한다. 단계 3: HMM은 L-module(고해상도, episodic)과 H-module(저주기, semantic)을 결합해 두 차원의 정보를 저장하고, HLR을 통해 이 두 표현을 점진적으로 변환한다. 단계 4: MemStack은 컨솔리데이션 경로의 중간 상태를 decoder의 초기화 정보로 다층에 걸쳐 전달해 긴 컨텍스트에서도 지속적으로 기억 정보를 활용한다.

두 모듈의 업데이트 주기를 다르게 설계하고, Hierarchical Latent Recursion(HLR)으로 l_state와 h_state를 교차 보정한다. 입력 토큰은 패치 단위로 xp로 변환되어 L-module과 H-module에 전달되며, 각 모듈은 내부 메모리 모듈과 게이팅된 어텐션으로 구성된다. 메모리 업데이트는 Mt = αMt−1 + St, St = ηtSt−1 − θt∇ℓ(Mt−1, xt)와 같은 형태의 업데이트 규칙으로 수행되며, ℓ은 associative 손실을 사용한다. 상호작용은 L_cycle과 H_cycle의 재귀를 통해 수행되며, 마지막에 H_state와 L_state를 Fusion MLP로 결합하여 출력으로 보낸다.

주요 벤치마크에서 Mela는 400M/800M/1.2B 파라미터 규모에서 Transformer++ 대비 개선된 Perplexity를 달성한다. 1024 토큰 컨텍스트에서 Mela 400M은 12.53, Transformer++ 400M은 13.59이다. 32768 컨텍스트에서 Mela 1.2B은 11.85, Transformer++ 1.2B은 597.37이다. 같은 4K 프리트레인 컨텍스트에서 Mela 계열은 12.01~9.08 수준의 PPL로 하락하며, 긴 컨텍스트에서도 성능 저하가 더 느리게 나타난다. Ablation 결과, H-cycle을 늘리면 긴 컨텍스트에서 이득이 커지고, L-cycle은 짧은 컨텍스트에서 큰 이득을 준다. 게이트를 제거하면 성능이 크게 떨어지며, ML P(MLP) 조합으로의 출력 결합이 가장 효과적이다. Newton-Schulz 정규화의 도입과 MemStack의 사용은 각각 PPL 개선에 기여한다.

아키텍처 개요: Mela의 HMM은 L-module(episodic memory)와 H-module(semantic memory)로 구성되며, xp 패치를 통해 입력을 패치 단위로 처리한다. 각 메모리 모듈은 memory block으로 이루어지고, Nc = Np/C로 정의된(chunk 수) 단위로 Chunk 내에서 어텐션과 메모리 모듈을 적용한다. Retrieval은 qt = Wq xt로 계산된 쿼리를 이용해 yt = N(qt; Mt−1)로 얻는다. 학습 목표로는 associative 손실을 사용하고, Mt 업데이트 규칙은 Moll−1의 gradient를 이용한 St를 통해 이루어진다. Hierarchical Latent Recursion(HLR)은 두 단계의 재귀로 구성되며, H_cycles와 L_cycles의 비율에 따라 L-module의 업데이트 빈도가 달라진다. with torch.no_grad() 구문을 사용하여 HLR의 상위 사이클은 gradient 없이 실행한 뒤, 마지막에 gradients를 적용한다. MemStack은 z(i) = Fusion(l(i), h(i))를 통해 각 cycle의 memory representation을 생성하고, α=1...L의 decoder layer에서 이 중간 memory representations를 주입한다. 최종 y(L)은 y(L−1)와 z(H) 및 z(i)를 순서대로 결합해 계산된다. 최적화는 Muon 옵티마이저를 사용하고 NewtonSchulz−k(St)로 St를 직교화한다. k=5로 설정하여 수렴 속도와 안정성을 확보한다. Patch 크기 P와 Chunk 크기 C, Memory blocks의 수는 모델 규모에 따라 조정되며, H-module은 더 많은 memory blocks를 가지도록 설계된다. MemStack은 중간 cycle의 memory representations를 디코더 초기층에 주입하므로 추가 토큰 없이도 memory 정보를 활용한다.