언어 모델은 수면이 필요한가? 온라인 추론 향상을 위한 오프라인 재현 기반 메커니즘

Transformer 기반 LLM의 self-attention은 컨텍스트 길이가 길어질수록 연산 복잡도가 증가한다. 본 연구는 컨텍스트 윈도가 가득 차면, 에 evicted된 컨텍스트를 고정된 크기의 fast weights로 합성하고, sleep 단계에서 N회에 걸친 오프라인 순환을 통해 fast weights를 업데이트한다. wake-time 예측은 한 번의 순방향 패스로 이루어지며, 더 깊은 추론을 sleep에 배분한다.

출발점과 한계: Self-attention은 과거 토큰 전체에 대한 관계를 계산해야 하므로 컨텍스트 길이가 길어질수록 연산이 기하급수적으로 증가한다. 빠른 가중치 메모리는 고정 크기의 상태에 요약된 정보를 저장하지만, 정보 손실이 발생해 깊은 순차 계산을 밀접하게 필요로 하는 태스크에서 한계가 나타난다. 논문은 evicted 컨텍스트를 fast weights로 변환하는 sleep-like offline recurrence를 제안한다. 해결 원리: 컨텍스트 윈도가 가득 차면, 모듈 내부의 SSM 블록에서 N회에 걸쳐 forward를 반복 수행하며 fast weights를 업데이트한다. 이때 계산은 학습 가능 질의 규칙(learned local rule)에 의해 진행되며, wake-time에는 KV 캐시를 버려도 추론은 단일 패스만으로 수행된다. 이 과정에서 gradient는 sleep 중의 재귀 업데이트를 통해 역전파된다. 달라지는 점: 추가적인 sleep 시간(N 회의 루프)을 통해 evicted 컨텍스트를 보다 풍부하게 정돈하고, 심층 순차 연산이 필요한 문제에서 성능 향상을 달성한다. 기존의 4-layer GDN-attention 하이브리드 모델은 t가 커질수록 성능 저하를 보이나, sleep를 도입하면 깊은 추론에 필요한 compute를 offline에 배분해 wake-time latency를 유지하면서 성능을 높인다.

Algorithm 1에 따라 학습이 진행된다. zero-initialize SSM fast weights S를 만들고, 토큰을 길이 L의 비중복 청크로 분할한다. consolidation phase에서는 mc가 모두 0인 경우 N번 루프를 수행하며(h,S를 갱신), prediction phase에서는 B_s sm 블록들을 통해 h와 S를 한 번 업데이트한다. 손실은 OutProj(h)와 실제 타깃 c, loss mask mc를 이용해 MaskedCE를 계산하고, 전체 그래프에 대해 역전파를 수행해 파라미터를 최적화한다. 이때 gradient는 sleep 동안의 fast weights에 흐른다. KV cache는 evicted 후 다음 청크로 넘어간다.

메인 벤치마크: cellular automaton에서 t 증가에 따라 vanilla 하이브리드가 무력해지는 반면, 2-loop은 약 20%, 3-loop/4-loop는 30% 이상 정확도 달성. Depo 태스크에서 4-loop 모델이 가장 어려운 16-hop에서 학습 속도와 정확도 개선을 보였다. GSM-Infinite에서 Jet-Nemotron 2B는 6-operation 문제에서 0.742→0.812로, 8-operation 문제에서 0.351→0.388로 개선되었고, Ouro 1.4B는 6-operation에서 0.419→0.615, 8-operation에서 0.210→0.272로 개선되었다. 슬라이딩-window eviction에서 op2는 52% 개선, op8은 약 30% 개선이 관찰되었다. 학습 처리량은 윈도우 크기가 충분히 크면 N의 증가에 따른 오버헤드 증가가 선형이지만, GPU 활용은 유지되었다. 이러한 실험은 sleep-like offline recurrence가 evicted 컨텍스트를 의미 있는 fast weights로 정리해 깊은 순차 추론을 가능하게 함을 보여준다.

아키텍처 구성: Embedding → Battn0 → … BattnD−1 → OutProj로 구성된 하이브리드/SSM 블록이 등장한다. eviction boundary에서 KV cache를 제거하기 전, N회의 반복 루프를 통해 fast weights S를 업데이트한다. S는 St = αtSt−1 + βt v_t k_t^⊤의 형태의 선형 재귀 업데이트 규칙을 따른다고 설명되며, ot = St qt로 표현된다. sleep 동안 그래디언트는 S로 흐르며, wake-time에는 단일 순방향 패스로 예측한다. 실험은 4-layer GDN-attention 하이브리드 구조를 기반으로 수행되었고, 다양한 N 값에서 성능 증가를 확인했다. Sliding-window eviction에서는 윈도우 L = 512, N ∈ {1,2,4}에서 특히 2/4 loop가 이점을 보인다. Throughput 분석은 window size가 커질수록 재귀 깊이에 따른 비용 증가가 선형적임을 보이고, 큰 L에서 Wall-time Utilization이 유지된다.