기술 요약: 메모리 내 추론 (Reasoning in Memory, RiM)

문제 제기

대형 언어 모델 (LLM) 의 추론 능력을 향상시키기 위한 기존 방법들은 최종 답변을 생성하기 전에 중간 토큰을 생성함으로써 테스트 시간 계산 (test-time compute) 을 확장하는 데 크게 의존합니다. 체인 오브 씽킹 (Chain-of-Thought, CoT) 추론이 그 예입니다. 그러나 이러한 접근법은 내부 계산을 자동회귀적 텍스트 생성에 직접적으로 결합합니다. 그 결과, 모델은 순수한 내부 추론이 아닌 문법적으로 유창한 자연어를 생성하는 데 상당한 계산 예산을 할당해야 하는, 즉 "소리를 내어 생각 (think out loud)"하도록 강요받게 됩니다.

최근의 "잠재적 추론 (latent reasoning)" 접근법들은 연속적인 표현을 사용하여 자연어 제약을 우회하려 시도하지만, 대부분 단계별 자동회귀적 생성 패러다임을 그대로 유지합니다. 이러한 방법들에서는 다음 단계가 이를 조건으로 삼기 전에 중간 계산들이 여전히 순차적으로 외부화됩니다. 본 논문은 인간의 인지 과정이 더 나은 설계 원리인 작업 기억 (working memory) 을 제공한다고 주장합니다. 인간은 모든 중간 생각을 외부화하지 않고도 작업과 관련된 정보를 내부적으로 보유하고 조작할 수 있습니다. 저자들은 LLM 이 토큰 생성의 순차적 병목 현상과 추론을 분리하는 유사한 내부 작업 공간을 활용하도록 훈련될 수 있다고 가정합니다.

방법론: 메모리 내 추론 (Reasoning in Memory, RiM)

저자들은 추론 단계의 자동회귀적 생성을 고정된 **메모리 블록 (memory blocks)**으로 대체하는 잠재적 추론 방법인 **메모리 내 추론 (Reasoning in Memory, RiM)**을 소개합니다.

핵심 메커니즘

추론 토큰의 시퀀스를 생성하는 대신, 모델은 $K$개의 고정된 메모리 블록 $m_{1:K}$를 입력으로 받습니다. 각 블록은 미리 정의된 특수 토큰 시퀀스로 구성됩니다:
$$m_k = [<b>, <m>, \dots, <m>, </b>]$$
여기서 $<b>$와 $</b>$는 블록 경계를 구분하고, $<m>$ 토큰들은 잠재적 작업 공간을 형성합니다.

• 고정 입력: 이러한 블록들은 고정된 토큰으로 입력 질문에 추가됩니다.
• 단일 순전파 (Single Forward Pass): 블록들이 생성된 출력이 아닌 고정된 입력이므로, 증강된 시퀀스 (질문 + 메모리 블록) 전체가 단일 순전파로 처리됩니다. 이는 자동회귀적 생성과 관련된 순차적 디코딩 병목 현상을 제거합니다.
• 학습된 표현: 토큰의 정체성과 위치는 고정되어 있지만, 그 맥락적 표현은 작업에 종속적으로 학습되어 중간 추론 단계를 인코딩합니다.

2 단계 커리큘럼

모델이 이러한 고정된 블록을 작업 기억으로 사용하도록 가르치기 위해 RiM 은 2 단계 훈련 커리큘럼을 사용합니다:

1 단계: 추론 단계 감독 (Reasoning Step Supervision)

• 목표: 메모리 블록에 계산적 역할을 부여하여 이를 기반 (ground) 으로 삼습니다.
• 메커니즘: $T$단계의 추론 궤적에 대해 모델은 $T$개의 메모리 블록으로 훈련됩니다. 각 블록 $m_t$ 이후, 모델은 다음 명시적 추론 단계 $r_{t+1}$를 예측해야 합니다.
• 제약: 사용자 정의 어텐션 마스크를 사용하여 단계 $t+1$의 읽기 출력 (readout) 이 질문과 이전 메모리 블록에는 주의를 기울일 수 있지만, 이전 명시적 추론 단계에는 주의를 기울일 수 없도록 합니다. 이는 모델이 필요한 중간 계산을 메모리 블록의 잠재적 작업 공간 내에 인코딩하도록 강제합니다.
• 감독: 점진적으로 감독 강도를 줄이는 어닐링 (annealing) 일정을 적용한 가중 음의 로그 가능도 (negative log-likelihood) 목적 함수를 사용합니다.

2 단계: 최종 답변 정제 (Final Answer Refinement)

• 목표: 기반을 둔 메모리 블록을 최종 답변 예측을 위한 잠재적 작업 공간으로 전환합니다.
• 메커니즘: 명시적 추론 궤적은 폐기됩니다. 모델은 고정된 수의 $K$개 메모리 블록을 사용합니다. 각 블록 이후, 모델은 최종 답변을 직접 예측하도록 훈련됩니다.
• 목적: 이는 중간 단계를 예측하는 것에서 더 많은 메모리 블록 (잠재적 계산) 이 사용 가능해짐에 따라 최종 답변을 점진적으로 정제하는 것으로 감독을 전환합니다.
• 추론: 테스트 시, 모델은 질문과 메모리 블록을 한 번에 처리하며, 마지막 블록의 읽기 출력이 답변을 제공합니다.

어텐션 마스킹

사용자 정의 인과적 어텐션 마스크 (논문 내 Figure 2) 는 이 방법의 핵심입니다. 이는 다음을 보장함으로써 정보 누출을 방지합니다:

1. 메모리 블록은 질문과 이전 메모리 블록에 주의를 기울입니다.
2. 대상 읽기 출력 (1 단계의 추론 단계, 2 단계의 최종 답변) 은 질문과 사용 가능한 메모리 블록에 주의를 기울이지만, 다른 대상 읽기 출력에는 주의를 기울이지 않습니다.
   이를 통해 모델이 미래의 정답 (ground truth) 을 읽음으로써 "속임수"를 쓰지 않고도 모든 대상이 단일 순전파로 예측될 수 있습니다.

주요 기여

1. 메모리 내 추론 (RiM): 중간 계산을 위한 전용 고정 잠재적 작업 공간을 제공하여 자동회귀적 생성과 이를 분리함으로써 LLM 의 작업 기억 용량을 해제하는 새로운 잠재적 추론 방법.
2. 2 단계 커리큘럼: 명시적 단계 감독을 통해 메모리 블록을 기반으로 삼은 후 최종 답변 예측을 위해 이를 정제하는 구조화된 훈련 접근법으로, 복잡한 부가적 목적 함수 없이 잠재 공간의 효과적인 사용을 가능하게 함.
3. 효율성과 성능: 기존 명시적 및 잠재적 추론 기준선과 동등하거나 더 나은 성능을 달성하면서도 직접 답변 모델의 추론 속도 (순차적 생성 오버헤드 없음) 를 유지한다는 실증적 증거.

실험 결과

저자들은 GPT-2와 Llama-3.2(1B 및 3B 파라미터) 모델을 사용하여 GSM8K(동일 분포) 및 GSM-Hard(이질 분포) 수학 추론 벤치마크에서 RiM 을 평가했습니다.

• 성능: RiM 은 가장 강력한 잠재적 추론 기준선 (Coconut) 과 직접 답변 감독 미세 조정 (SFT) 을 일관되게 능가했습니다.
  • GSM8K(Llama-3.2-1B) 에서 RiM 은 **42.1%**의 탐욕적 정확도 (greedy accuracy) 를 달성했으며, Coconut 은 36.9%, CoT 없는 SFT 는 **23.9%**였습니다.
  • GSM-Hard 에서 RiM 은 **10.5%**를 달성하여 Coconut(8.5%) 과 CoT 없는 SFT(5.3%) 를 능가했습니다.
  • RiM 은 또한 Pass@8 점수에서 우월성을 보여, 올바른 솔루션이 높은 확률로 잠재적 궤적에 분포되어 있음을 나타냈습니다.
• 지연 시간: RiM 은 CoT 없는 SFT 와 비교 가능한 추론 속도를 달성했습니다 (Llama-3.2-1B 의 경우 첫 번째 토큰 도달 시간 약 16ms). 반면 Coconut 은 약 7 배, CoT 가 있는 SFT 는 자동회귀적 생성으로 인해 약 27 배 느렸습니다.
• 표현 분석: 메모리 블록 표현의 PCA 투영은 모델이 고유하고 샘플에 종속적인 잠재적 작업 공간을 생성하도록 학습함을 보여주었습니다. 표현들은 훈련 중에 체계적으로 진화하여, 블록들이 단순한 고정 자리 표시자가 아닌 활성 계산 상태임을 확인시켰습니다.
• 강건성: 성능은 다양한 추론 시간 메모리 예산 (블록 수) 에서 안정적으로 유지되어, 이 방법이 배포 제약 조건에 강건함을 시사합니다.

중요성과 주장

본 논문은 LLM 이 잠재적 추론을 위한 효과적인 메커니즘으로 작업 기억을 사용하도록 훈련될 수 있다고 주장합니다.

주요 중요성은 계산과 의사소통의 분리에 있습니다. 중간 생각의 생성을 고정된 메모리 블록으로 대체함으로써 RiM 은 모델이 자연어 생성이나 순차적 연속 표현의 오버헤드 없이 밀집되고 병렬화 가능한 내부 계산을 수행할 수 있게 합니다. 이는 다음과 같은 방법을 산출합니다:

1. 계산적 효율성: 자동회귀적 생성의 순차적 병목 현상을 피하여 단일 통과 추론을 가능하게 함.
2. 효과성: 수학 벤치마크에서 기존 잠재적 추론 방법들 중 최첨단 성능 달성.
3. 확장성: 복잡한 순환 구조나 다중 경로 훈련 목적 함수를 요구하지 않고 다양한 모델 계열과 크기에서 작동함.

저자들은 이 접근법이 텍스트로 외부화하는 대신 추론 과정을 내부화함으로써 인간의 인지 효율성을 모방하는 더 효율적인 추론 시스템으로 가는 실현 가능한 길을 보여준다고 결론지었습니다.