Memoir: Should a Model Write to Its Memory While It Thinks?

기술 요약: Memoir – 모델이 생각하는 동안 자신의 메모리에 기록해야 하는가?

문제 정의
현재의 언어 모델은 일반적으로 순전파(forward pass) 과정 동안 훈련된 파라미터를 고정된 상태로 유지하며, 내부 가중치 상태를 변경하지 않고도 유효한 메모리를 확장하기 위해 외부 검색이나 긴 컨텍스트에 의존한다. 테스트 타임 훈련(test-time training)과 적응형 재귀(adaptive recurrence)를 통해 모델이 은닉 상태를 업데이트하거나 어려운 입력에 더 많은 계산량을 할당할 수는 있지만, 특정 기술적 위험이 여전히 해결되지 않은 채 남아 있다: 바로 모델이 자신의 빠른 메모리(fast memory)를 읽고 쓰는 과정을 동시에 수행하는 재귀적 잠재 프로세스(recurrent latent process)의 안정성 문제이다.

본 연구에서 조사하는 핵심 질문은 모델이 매 폰더링(pondering, 재귀적) 반복 단계 동안 자신의 빠른 메모리 계층을 다시 쓸 수 있도록 허용해야 하는가 하는 점이다. 저자들은 특정한 실패 모드를 가설로 세웠다: 만약 잠재 표현(latent)이 정교해지는 동안 메모리가 다시 쓰인다면, 에너지 헤드(energy head)가 잠재 표현 자체가 더 일관되게 변해서가 아니라, 단지 현재의 잠재 표현과 일치하도록 메모리가 이동했기 때문에 개선된 것처럼 보일 수 있다는 것이다. 이러한 "결합된(coupled)" 쓰기 방식은 계산 루프 중에 비정상 상태(non-stationary)의 연산자를 생성하여 학습 신호를 오염시킬 가능성이 있다.

방법론
본 논문은 네 가지 가중치 계층, 가변 깊이의 잠재 재귀, 그리고 미래 잠재 에너지 목적 함수를 통합한 Memoir 프레임워크를 소개한다.

1. 4단계 메모리 인터페이스:

   • $W_0$ (Fast): 매 순전파 시마다 쓰기가 가능한 샘플별 행렬.
   • $W_1$ (Episodic): 세그먼트당 학습된 게이트를 통해 업데이트되는 샘플별 행렬.
   • $W_2$ (Slow): 드물게 업데이트되는 공유 훈련 가능 파라미터(consolidation).
   • $W_3$ (Frozen): 업데이트되지 않는 공유 파라미터.
     이 인터페이스는 소유권과 쓰기 경계를 명확히 정의하여, 샘플별 가소성 상태와 공유되는 느린/고정된 상태를 구분한다.

2. 가변 깊이 폰더링(Variable-Depth Ppering):
   모델은 공유 블록이 현재의 잠재 표현과 사용 가능한 계층을 읽어 새로운 상태($h^{(\ell+1)}$)를 제안하는 재귀적 잠재 프로세스를 채택한다. 중단 헤드(halting head)는 반복을 언제 멈출지를 결정한다.

   • 결합된 팔 (Arm A - Coupled Arm): 활성 폰더링 반복 중에 모델은 빠른 계층($W_0$)을 읽고, 현재의 잠재 표현을 기반으로 한 델타 규칙(delta-rule) 업데이트를 사용하여 즉시 $W_0$에 기록한다.
   • 읽기 전용 팔 (Arm B - Read-Only Arm): 모델은 폰더링 중에 빠른 계층을 읽지만, 순전파 경계에 도달할 때까지는 $W_0$에 기록하지 않는다. 이 경우 루프 동안 연산자는 정상 상태(stationary)를 유지한다.

3. 목적 함수:

   • 미래 잠재 예측 (Future-Latent Prediction): 현재 표현으로부터 타겟 잠재 표현 $z_{t+\Delta}$를 예측하며, 이는 stop-gradient EMA(지수 이동 평균) 텔러(teacher)를 대상으로 훈련된다.
   • 에너지 목적 함수 (Energy Objective): 에너지 헤드는 표현 쌍의 점수를 매긴다. 마진(margin)은 $m = E[E(\text{corrupted})] - E[E(\text{coherent})]$로 정의된다. 양수의 마진은 일관된 쌍이 더 낮은 에너지를 받는다는 것을 의미한다.

4. 실험 설계:
   연구는 키 간섭(key interference)이 포함된 절차적 연상 회상(procedural associative recall) 과업(64개 값, 16개 방해 요소, 간섭 0.85)을 사용한다.

   • 대조군 (Control): 두 팔(arm)은 동일한 파라미터 수(총 81,738개, 훈련 가능 76,362개), 동일한 순전파 곱셈-누산(multiply-accumulate) 횟수, 데이터 및 시드를 가진다.
   • 예산 (Budgets): 주요 테스트는 240 스텝(12개 시드)에서 진행되며, 장기 대조군은 960 스텝(8개 시드)에서 진행된다.
   • 구현: 본 시스템은 루프 내 쓰기 존재 여부만을 차이점으로 하는 매칭된 아키텍처 절제 실험(ablation)을 포함한다.

주요 결과

1. 학습 속도 저하 (Learning Speed Penalty): 고정된 예산인 240 스텝에서, **읽기 전용 팔 (Arm B)**이 **결합된 팔 (Arm A)**보다 유의미하게 높은 성능을 보였다.

   • 읽기 전용 회상 (Read-Only Recall): 0.6557 (95% CI: [0.5953, 0.7160])
   • 결합된 회상 (Coupled Recall): 0.5203 (95% CI: [0.4522, 0.5883])
   • 차이: 0.1354. 12개 시드에 대한 대응 표본 t-검정 결과 $t=3.23$ (df=11)이었으며, 차이에 대한 95% 신뢰 구간은 [0.0431, 0.2277]로 0을 포함하지 않았다. 읽기 전용 팔이 12개 시드 중 10개에서 승리했다.

2. 수렴 동작 (Convergence Behavior): 960 스텝에서 두 팔 모두 과업의 천장(ceiling)인 1.0000에 도달했으며, 성능 격차는 사라졌다. 이는 결합된 아키텍처가 이 특정 과업에서 더 낮은 도달 가능한 능력을 갖는 것이 아니라, 측정된 예산 하에서 학습 속도 저하를 겪는다는 것을 나타낸다.

3. 에너지 마진 안정성 (Energy Margin Stability): 메모리 재작성이 현재의 잠재 표현과 일치하도록 만들어 에너지 신호를 오염시킨다는 예측된 실패 모드는 발생하지 않았다. 에너지 마진은 훈련 과정에서 거의 0에서 0.0231까지 성장하였으며, 이는 결합된 쓰기가 이 설정에서 에너지 목적 함수를 불안정하게 만들지 않았음을 입증한다.

4. 엔지니어링 개선 사항: 커널 구조 조정을 통해 델타 규칙 순전파 시간을 0.907 ms에서 0.351 ms로 단축(2.59배 속도 향상)하였으며, NVIDIA GeForce RTX 5090 Laptop GPU에서 유효 대역폭을 18.6 GB/s에서 48.0 GB/s로 증가시켰다.

의의 및 주장
본 논문은 다섯 가지 구체적이고 검증 가능한 기여를 한다:

1. 인터페이스 명세: 명시적인 소유권, 수명 및 쓰기 경계를 가진 4단계 메모리 인터페이스.
2. 프로세스 정의: 주변 모델을 변경하지 않고도 빠른 계층의 쓰기 정책을 전환할 수 있는 수직적 가변 깊이 폰더링 프로세스.
3. 엄격한 절제 실험: 240 스텝에서의 12개 시드 기반 절제 실험과 960 스텝에서의 8개 시드 대조 실험을 통해 루프 내 쓰기의 효과를 분리하여 검증함.
4. 부정적 결과 검증: 예측된 에너지 마진의 붕괴가 발생하지 않음을 확인하여, 결합된 폰더링이 본 구성에서 에너지 신호를 본질적으로 오염시킨다는 가설을 반박함.
5. 구현 투명성: 디바이스 수준의 타이밍 및 메모리 측정을 포함한 전체 코드 공개 (memoir/ 내 5,920행, 테스트 코드 2,061행).

한계 및 제한적 주장
저자들은 실험 결과가 결합된 아키텍처가 더 낮은 궁극적 능력을 갖는다는 것을 증명하는 것은 아니라고 명시적으로 밝히고 있다. 두 팔 모두 포화된 과업(saturated task)에서 천장에 도달했기 때문이다. 관찰된 격차는 능력의 저하가 아닌 학습 속도의 저하이다. 또한, 본 실험은 쓰기의 타이밍과 쓰기의 양을 분리하지 않았다. 결합된 팔은 읽기 전용 팔보다 에피소드당 약 7배 더 많은 쓰기 작업을 수행한다. 또한 본 연구는 자연어, 장기 호흡의 지속적 학습, 파괴적 망각 저항성 또는 학습된 업데이트 규칙을 테스트하지 않았다. 에너지 마진은 환각(hallucination)에 대한 일반적인 검증 도구가 아니라 과업 특화된 진단 도구로 보고되었다.

결론적으로, Memoir는 추론 중 쓰기 가능한 메모리를 테스트하기 위한 구체적인 인터페이스를 제공한다. 주요 발견은 특정 절차적 회상 과업과 훈련 예산 하에서, 폰더링 중에 빠른 메모리를 다시 쓰도록 허용하는 것이 읽기 전용 방식에 비해 학습 속도를 늦추지만, 결국 수렴을 막지는 못한다는 것이다.