Language Models Need Sleep

기술 요약: "언어 모델은 수면이 필요하다"

문제 제기
트랜스포머 기반 대규모 언어 모델 (LLM) 은 어텐션 메커니즘이 컨텍스트 길이에 따라 이차적으로 확장되고 키-값 (KV) 캐시 메모리가 선형적으로 증가하기 때문에 장기 과제를 수행하는 데 어려움을 겪습니다. 최근 등장한 효율적인 하이브리드 모델 (어텐션과 상태 공간 모델, SSM 을 결합) 은 고정 크기의 "빠른 가중치 (fast weights)"를 통해 메모리 비용을 완화하지만, 추론 깊이에서 치명적인 병목 현상에 직면합니다.

저자들은 순수 SSM-어텐션 하이브리드 모델이 KV 캐시에서 제거된 토큰에 대한 깊은 순차 계산을 수행하지 못함을 입증했습니다. 메모리 용량이 정보를 저장하기에 충분하더라도, 요구되는 추론 깊이가 증가함에 따라 (예: 많은 단계에 걸친 세포 자동자 시뮬레이션 또는 다중 홉 그래프 탐색 수행) 모델의 성능이 저하됩니다. 이 병목 현상은 단순히 저장 용량의 문제가 아니라, 단일 순전파 (forward pass) 내에서 제거된 컨텍스트를 유용한 내부 상태로 변환하기에 충분한 계산 자원이 부족하다는 점에 기인합니다.

방법론: LLM 수면
생물학적 기억 고정 (수면 중 해마 재생) 에서 영감을 받아, 이 논문은 LLM 을 위한 "수면과 유사한" 메커니즘을 제안합니다. 이 접근법은 추가 계산을 오프라인 단계로 이동시켜 온라인 예측 단계의 지연 시간을 유지합니다.

1. 아키텍처: 모델은 빠른 가중치 메모리 ($S_t$) 와 함께 어텐션 블록 ($B^{attn}$) 과 SSM 블록 ($B^{ssm}$) 을 번갈아 배치한 하이브리드 아키텍처를 사용합니다.
2. 수면 단계: 컨텍스트 창이 가득 차면 (예: 매 $L$개 토큰마다), 모델은 "수면" 단계에 진입합니다. 이 단계 동안:
   • 모델은 외부 입력 토큰을 받지 않습니다.
   • 누적된 컨텍스트에 대해 **$N$회의 오프라인 순환 통과 (recurrent passes)**를 수행합니다.
   • 각 통과에서 모델은 학습된 로컬 규칙 (예: 게이트된 델타 규칙 업데이트) 을 사용하여 SSM 블록 내의 빠른 가중치를 재귀적으로 업데이트합니다.
   • 이 과정을 통해 모델은 제거된 컨텍스트의 압축된 표현을 반복적으로 정제할 수 있습니다.
3. 제거 및 예측: $N$회의 통과 후 KV 캐시가 비워지거나 (슬라이딩 윈도우 변형의 경우 부분적으로 비워짐), 모델은 업데이트된 빠른 가중치를 사용하여 작동을 재개합니다.
4. 추론 제약: 결정적으로 실제 예측 단계 (질문 응답) 동안 모델은 단일 순전파만 수행합니다. 모든 깊은 추론과 컨텍스트 고정은 이전 "수면" 단계 동안 완료되어야 합니다.
5. 학습: 모델은 컨텍스트가 제거된 후 예측 오차를 최소화하기 위해 순환 수면 통과를 포함한 전체 계산 그래프를 통해 역전파로 종단 간 (end-to-end) 학습됩니다.

주요 기여

• 추론 병목 현상 식별: 저자들은 통제된 환경에서 순수 SSM-어텐션 하이브리드가 메모리 용량이 고정되어 충분함에도 불구하고 깊은 추론이 필요한 작업 (세포 자동자의 높은 $t$, 그래프 검색의 높은 $k$) 에서 실패함을 보여줍니다. 이는 단순히 메모리가 아닌 계산 깊이를 제한 요인으로 규명합니다.
• 수면 아키텍처: 기억 고정을 위해 순환 계산을 통합하는 아키텍처를 제안합니다. 수면 통과 횟수 ($N$) 를 늘림으로써 모델은 추론 지연 시간을 증가시키지 않고도 더 깊은 추론 작업을 처리할 수 있도록 계산 능력을 확장할 수 있습니다.
• 실증적 검증: 이 방법은 다음에서 검증되었습니다:
  • 룰 110 (세포 자동자): 깊은 순차 시뮬레이션이 필요한 작업.
  • Depo: 조각난 컨텍스트에 대한 그래프 탐색이 필요한 다중 홉 지식 검색 작업.
  • GSM-Infinite: 긴 컨텍스트와 방해 토큰을 포함하는 현실적인 수학 추론 벤치마크로, 사전 학습된 LLM(Jet-Nemotron 2B 및 Ouro 1.4B) 에서 평가됨.

결과

• 합성 작업: 수면 기간 ($N$) 을 늘리면 성능이 크게 향상됩니다. $t=32$인 룰 110 작업에서 순환이 없는 모델 ($N=1$) 은 무작위 추측 수준 (10% 정확도) 에 근접한 성능을 보인 반면, 34회의 수면 통과를 가진 모델은 30% 이상의 정확도를 달성했습니다. 마찬가지로 Depo 작업에서 $N$을 늘리면 학습 속도가 빨라지고 정확도가 향상되었으며, 특히 단일 통과 모델이 정체된 고홉 쿼리 (예: 16 홉) 에서 두드러졌습니다.
• 현실적 추론 (GSM-Infinite): 사전 학습된 모델에서 $N$을 늘리면 더 어려운 문제 (6~8 개의 산술 연산이 필요한 경우) 에서 상당한 향상을 보였습니다. 예를 들어, Ouro 1.4B 모델에서 수면 통과를 1 회에서 4 회로 늘리면 6 개 연산 문제의 정확도가 41.9% 에서 61.5% 로 향상되었습니다.
• 슬라이딩 윈도우: 모델이 최근 토큰을 유지하더라도 이전 컨텍스트를 압축하기 위해 수면에 의존하는 슬라이딩 윈도우 제거 전략에서도 이러한 이점이 지속되었습니다.
• 학습 처리량: 수면 통과의 순차적 특성으로 인해 $N$에 따라 학습 처리량이 거의 선형적으로 감소하지만, 저자들은 큰 컨텍스트 창의 경우 GPU 가 여전히 포화 상태이므로 완전히 병렬화된 기준선과 비교할 때 실제 벽시계 학습 시간을 크게 방해하지 않는다고 지적합니다.

의의 및 주장
이 논문은 오프라인 순환이 제거된 컨텍스트를 빠른 가중치로 조직화하는 실현 가능한 메커니즘이라고 주장하며, 모델이 더 이상 직접 주시할 수 없는 정보에 대해 깊이 추론할 수 있게 한다고 합니다.

• 적은 범위: 저자들은 연구 결과가 통제된 합성 작업과 소규모 사전 학습 모델을 기반으로 함을 강조합니다. 모든 긴 컨텍스트 추론 문제를 해결한다고 주장하지는 않지만, "수면"(오프라인 고정) 이 메모리 용량과 추론 깊이 사이의 간극을 메우는 구체적인 아키텍처 경로를 보여줍니다.
• 생물학적 영감: 이 연구는 단기 컨텍스트를 장기 가중치로 전환하려면 반복적 계산 (재생) 이 필요하며, 이는 실시간 예측의 엄격한 지연 시간 제약과 분리될 수 있다는 가설을 검증합니다.
• 향후 방향: 이 논문은 이를 더 능력 있는 긴 컨텍스트 시스템을 구축하기 위한 방법론적 기여로 위치시키며, 추론 시간의 추론 깊이를 위해 학습 시간의 순차성을 교환하는 것이 향후 연구의 유망한 방향임을 시사합니다.