Unlocking the Working Memory of Large Language Models for Latent Reasoning
https://arxiv.org/abs/2605.30343
Lukas Aichberger, Sepp Hochreiter | ELLIS Unit Linz and LIT AI Lab, Institute for Machine Learning, Johannes Kepler University Linz; NXAI GmbH | arXiv:2605.30343 | 2026년 5월
1. 서론: 자연어 CoT 비용과 내부 작업 기억의 문제
Reasoning in Memory(RiM)가 겨냥하는 병목은 단순히 더 많은 추론 토큰을 쓰느냐의 문제가 아니다. 최근 언어 모델의 수학·코딩·계획 성능은 chain-of-thought, self-consistency, test-time search, process verifier처럼 중간 추론을 길게 외부화하는 방식으로 많이 끌어올려졌다. 그러나 이 흐름은 중간 계산을 언어 문장으로 내보내고, 그 문장을 다시 다음 계산의 입력으로 삼는 구조에 묶여 있다. 모델이 실제로 필요한 것은 문제 상태, 부분 계산, 가능한 답의 후보, 실패한 가지를 조작하는 내부 표현일 수 있는데, 표면적으로는 문법적으로 자연스러운 설명문을 계속 생성해야 한다.
이 논문은 그 비용을 communication과 computation의 결합으로 설명한다. 자연어는 사람에게 설명하기 좋은 매체지만, 모델 내부에서 모든 중간 상태를 다루기 위한 최적의 계산 매체라고 보기는 어렵다. CoT는 계산을 보이게 만든다는 장점이 있으나, 매 단계가 token decoding으로 진행되기 때문에 latency와 메모리 사용량이 늘고, 문장 형식에 맞추는 부담도 함께 생긴다. 특히 작은 모델에서는 reasoning trace를 길게 생성하는 동안 오류가 누적될 수 있고, 큰 모델에서는 긴 추론이 비용 상한을 빠르게 소진한다.
기존 latent reasoning 연구는 이 문제를 알고 있었다. Coconut처럼 연속 표현을 다시 입력으로 넣는 방법은 자연어 문장 대신 hidden representation을 중간 상태로 쓰며, pause token이나 filler token 계열 방법도 의미가 고정되지 않은 토큰에 계산 여유를 주려 한다. 하지만 논문이 지적하는 핵심은, 이런 방법 중 상당수가 여전히 순차 생성 패러다임 안에 남아 있다는 점이다. discrete token이든 continuous thought이든, 다음 계산이 이전 중간 표현에 의존하려면 그 표현을 먼저 생성해 sequence에 붙여야 한다.
RiM은 이 결합을 다른 방식으로 끊는다. 모델이 중간 추론 표현을 한 단계씩 내보내게 하는 대신, 입력 질문 뒤에 고정된 memory block을 붙이고 이 블록의 contextual representation을 내부 작업공간으로 쓰게 한다. token identity와 위치는 고정되어 있지만, 각 block의 hidden state는 질문과 앞선 block에 따라 달라진다. 따라서 모델은 외부로 reasoning step을 생성하지 않아도, forward pass 안에서 문제별 latent workspace를 구성할 수 있다.
이 관점은 작업 기억에 대한 인지과학적 비유를 가져온다. 사람도 복잡한 계산을 할 때 항상 모든 중간 생각을 말로 출력하지 않는다. 필요한 변수, 부분 결과, 조건을 머릿속 작업공간에 유지하고 조작한 뒤 최종 답만 내놓을 수 있다. RiM은 이 비유를 LLM 입력 형식으로 내린다. 즉, 질문과 답 사이에 문장형 reasoning trace를 길게 놓는 대신, 질문 뒤에 fixed special token blocks를 놓고 그 hidden state를 계산 공간으로 훈련한다.
논문이 흥미로운 이유는 성능 주장이 과장된 범용 지능 선언이 아니라, 매우 구체적인 구조적 trade-off로 제시된다는 점이다. Stage 1은 memory block이 다음 reasoning step을 복구하도록 훈련해 block을 의미 있는 내부 상태로 접지한다. Stage 2는 step supervision을 제거하고 각 block 뒤에서 최종 답을 예측하게 해 답을 점진적으로 정제한다. 그 결과 RiM은 CoT처럼 긴 문장을 생성하지 않으면서도, Coconut 같은 latent baseline과 비교 가능한 정확도를 보이고, Llama-3.2-1B 기준 latency는 직접 답변 SFT와 거의 같은 수준으로 유지된다.
Figure 1: RiM의 두 단계 훈련 개요: memory block을 reasoning step 예측으로 접지한 뒤 final answer refinement로 전환한다.
Figure 1은 RiM의 핵심 흐름을 한 장에 압축한다. 노란색 memory block은 입력 뒤에 고정된 special token 묶음으로 들어가고, Stage 1에서는 각 block 뒤 readout이 다음 reasoning step을 예측하도록 학습된다. Stage 2에서는 그 중간 step 목표를 제거하고, 각 block이 최종 답을 점점 더 잘 예측하도록 훈련 신호를 바꾼다. 이 설계 덕분에 추론 표현은 외부 문장으로 생성되지 않고 내부 hidden state로 남는다.
2. 배경 및 관련 연구: CoT, latent reasoning, filler token의 한계
2.1 CoT와 test-time compute의 비용 구조
Chain-of-thought(CoT)는 LLM reasoning 연구의 중요한 출발점이다. 중간 풀이를 생성하게 하면 모델이 복잡한 문제를 여러 작은 단계로 분해하고, self-consistency나 verifier를 붙이면 여러 후보 trace 중 더 나은 답을 고를 수 있다. DeepSeek-R1 계열처럼 RL로 장문 추론 행동을 끌어올리는 흐름도 같은 축에 놓인다. 그러나 CoT의 계산은 결국 autoregressive text generation이다. 모델은 단계마다 다음 토큰을 예측하고, 그 토큰이 다시 다음 단계의 조건으로 들어간다.
이 구조는 해석 가능성과 비용을 동시에 낳는다. 중간 trace가 보이면 사람이 오류를 추적하거나 verifier가 step을 평가하기 쉽다. 반대로 모든 중간 계산이 언어로 외부화되므로, latency는 생성 token 수에 거의 비례하고, 실제 계산과 설명을 위한 문장 생산이 뒤섞인다. 논문은 이 점을 LLM이 생각을 말하는 방식에 계속 묶여 있는 상태로 본다. RiM은 중간 상태를 완전히 숨기는 대신, 훈련 중에는 step supervision으로 접지하고 추론 시에는 fixed memory block으로 병렬 처리하는 절충안을 제시한다.
2.2 latent reasoning이 남긴 미해결 지점
Coconut과 같은 explicit latent reasoning은 자연어 token 대신 continuous representation을 sequence에 넣는다. 이는 reasoning이 반드시 문장이어야 한다는 전제를 약하게 만든다. 하지만 representation을 한 칸씩 만들어 다시 입력으로 넣는 구조라면, 시간축은 여전히 순차적이다. 중간 표현의 표면은 바뀌었지만, 다음 계산이 이전 표현을 기다리는 병목은 남는다. 논문은 이 차이를 강조하면서, RiM이 continuous thought 자체보다 fixed workspace라는 실행 방식을 택한다고 설명한다.
filler token 계열 연구도 같은 문제를 다른 방식으로 다뤄 왔다. 의미 없는 토큰을 추가하면 모델이 그 위치를 계산용 scratchpad처럼 쓸 수 있을 것 같지만, 실제로는 그냥 무시되거나 성능을 떨어뜨릴 수 있다. 따라서 논문은 memory block이 유용해지려면 단순 삽입보다 더 강한 훈련 신호가 필요하다고 본다. Stage 1의 reasoning step supervision은 바로 이 요구에 대응한다. block 뒤의 readout이 다음 풀이를 맞히려면, block 내부 representation은 이전 정보와 현재 부분 계산을 담아야 한다.
2.3 기존 wiki 흐름과의 연결
이전에 리뷰한 Path-Lock Expert는 하이브리드 사고 모델에서 reasoning mode가 직접 답변 모드로 새는 reasoning leakage를 다뤘다. 거기서 핵심은 어떤 출력 모드가 선택됐는지, 그리고 MLP 경로가 그 선택을 얼마나 잘 유지하는지였다. RiM은 출력 모드보다 더 앞쪽의 질문을 던진다. 중간 추론을 꼭 출력 경로로 흘려보내야 하는지, 아니면 fixed memory block 내부에서 처리하고 답만 읽어낼 수 있는지를 실험한다.
또 다른 연결점은 Memory Grafting이다. 그 논문은 강한 grafting model의 hidden state를 frozen latent memory bank로 저장하고, n-gram lookup을 통해 recipient model에 접목한다. RiM 역시 latent memory라는 표현을 쓰지만, 외부 검색 테이블을 붙이는 구조와는 다르다. RiM의 memory block은 질문마다 새로 contextualized되는 내부 작업공간이며, retrieval key나 offline teacher bank 없이 동일한 forward pass 안에서 구성된다. 따라서 두 논문은 모두 언어 모델 안팎의 memory capacity를 다루지만, 하나는 외부 지식·표현을 저장하고, 다른 하나는 추론 도중의 작업 상태를 담는다.
| Table 1: RiM이 기존 reasoning 방식과 갈라지는 지점 | |||
|---|---|---|---|
| 축 | CoT / SFT w CoT | Coconut 계열 | RiM |
| 중간 표현 | 자연어 reasoning trace | continuous thought를 순차 생성 | fixed special token block의 hidden state |
| 실행 방식 | token을 길게 생성 | latent representation도 단계별 생성 | memory block을 한 번의 forward pass로 처리 |
| 훈련 신호 | 정답 또는 CoT trace next-token prediction | curriculum 및 continuous feedback | Stage 1 step supervision, Stage 2 final answer refinement |
| 장점 | 보이는 trace와 verifier 결합이 쉬움 | 언어 표면 비용 일부 완화 | 직접 답변 latency에 가까운 latent workspace |
| 주요 위험 | 긴 trace 비용과 오류 누적 | 순차 feedback 병목 유지 | memory block 의미가 과제 밖에서 유지되는지 추가 검증 필요 |
Table 1은 RiM을 단순한 CoT 압축 기법으로 보지 않게 해 준다. 논문은 자연어 trace를 짧게 요약하거나 continuous token으로 치환하는 데서 멈추지 않고, 중간 계산의 배치 자체를 바꾼다. 이 차이 때문에 RiM의 성능 비교는 accuracy만 보면 부족하고, time-to-first-token, pass@k, block별 answer transition, latent representation 구조를 함께 봐야 한다.
3. 방법론: Reasoning in Memory(RiM)
3.1 fixed memory block의 입력 형식
RiM은 질문 token sequence 뒤에 $K$개의 memory block을 붙인다. 각 block은 시작 delimiter, 여러 개의 memory token, 종료 delimiter로 구성된다. 논문은 기본적으로 block 내부 memory token 수를 $M=2$로 둔다. 중요한 점은 이 token들이 미리 정해진 surface meaning을 갖지 않는다는 것이다. <m> token은 사람이 읽을 수 있는 풀이 문장이 아니라, 모델이 contextual hidden state를 만들기 위한 자리다.
저자들은 기존 vocabulary embedding을 모두 바꾸기보다, 새로 추가한 special token embedding만 업데이트한다. 이 선택은 pretrained language knowledge를 불필요하게 흔들지 않기 위한 안전장치로 읽을 수 있다. memory block이 성능을 높인다면 그것은 기존 단어 embedding 전체가 새로 적응했기 때문이 아니라, 새 special token 위치가 질문별 작업 상태를 담도록 학습됐기 때문이어야 한다. 이 격리는 작은 모델과 큰 모델을 비교할 때도 중요하다.
형식적으로 RiM은 질문 $\mathrm{x}$, 답 $\mathrm{y}$, memory block $\mathrm{m}_{1:K}$를 함께 사용한다. 기존 CoT가 $\mathrm{x}\rightarrow \mathrm{r}_{1:T}\rightarrow \mathrm{y}$의 긴 token 경로를 만든다면, RiM은 $\mathrm{x}\rightarrow \mathrm{m}_{1:K}\rightarrow \mathrm{y}$의 형태로 중간 추론 공간을 고정한다. 각 $\mathrm{m}_k$의 token identity는 같아도 hidden representation은 질문과 이전 block에 의존해 달라지므로, block은 비어 있는 슬롯을 넘어 input-conditioned latent state로 작동한다.
Figure 2: RiM attention mask: memory block과 target reasoning step 사이의 정보 흐름을 제한해 병렬 step supervision을 가능하게 한다.
Figure 2는 RiM이 Stage 1에서 dense supervision을 넣는 방식을 보여준다. memory block은 질문과 이전 block을 볼 수 있고, 각 target reasoning step은 자신보다 앞선 memory block을 읽지만 다른 target step은 보지 못한다. 이 mask를 쓰면 여러 step target을 한 번의 forward pass에서 예측하면서도, 서로의 정답을 훔쳐보는 leakage를 막을 수 있다.
3.2 Stage 1: reasoning step supervision으로 block을 접지하기
memory block은 처음부터 작업 기억이 아니다. 아무 훈련 신호 없이 special token만 추가하면 모델은 그 위치를 무시하거나, 단순 위치 신호로만 사용할 수 있다. Stage 1의 목표는 block을 다음 reasoning step을 복구하는 데 필요한 내부 상태로 만드는 것이다. 각 memory block 뒤에서 readout을 붙이고, 해당 readout이 다음 reasoning step을 예측하도록 next-token prediction을 수행한다.
이때 readout이 볼 수 있는 정보가 제한된다. 특정 step의 target은 이전 memory block까지의 latent state를 보고 다음 풀이 조각을 맞혀야 한다. 따라서 모델은 memory block 안에 문제의 현재 상태, 이전 계산 결과, 다음 step을 준비하는 단서를 저장해야 한다. 논문은 이 구조를 JEPA류 predictive representation learning과도 연결한다. 보이지 않는 reasoning structure를 예측하게 함으로써 memory block이 task-relevant representation을 배우도록 만든다는 설명이다.
Stage 1의 손실은 각 reasoning step에 가중치를 둔 next-token loss로 구성된다. 여기서 중요한 것은 objective가 특수한 RL이나 별도 auxiliary head를 요구하지 않는다는 점이다. 표준 언어 모델 학습 절차를 유지하되, target 배치와 attention mask를 바꿔 dense supervision을 만든다. 따라서 RiM은 구현 관점에서 큰 구조 개조보다 입력 구성과 훈련 curriculum 설계에 가까운 방법이다.
Figure 3: RiM Stage 1과 Coconut curriculum 비교: RiM은 여러 단계의 progressive curriculum 대신 모든 memory block을 한 번에 supervision한다.
Figure 3은 RiM이 Coconut보다 curriculum을 단순화하는 지점을 보여준다. Coconut은 continuous thought 개수를 점진적으로 늘리는 여러 curriculum stage가 필요하지만, RiM은 fixed memory block 전체를 두고 reasoning step target을 병렬로 예측한다. 이 설계는 memory block마다 dense한 학습 신호를 주면서도, 추론 시에는 중간 thought를 생성하지 않는 방향으로 이어진다.
3.3 Stage 2: final answer refinement로 훈련 목표 전환
Stage 1이 memory block을 reasoning step에 맞춰 접지한다면, Stage 2는 그 block을 최종 답을 위한 작업공간으로 전환한다. 저자들은 step-level supervision을 제거하고, 각 memory block 뒤 readout이 final answer를 예측하도록 학습한다. 이때 block이 하나씩 늘어날수록 답이 더 정제될 수 있다. 즉, 모델은 첫 block만 보고도 답을 시도할 수 있고, 더 많은 block을 보면 앞선 내부 상태를 누적해 더 나은 final answer를 낼 수 있다.
이 구조는 any-block readout이라는 중요한 평가 모드를 가능하게 한다. 마지막 block만 읽는 final-block 방식은 가장 직접적인 latency 비교에 유리하다. 반면 any-block 방식은 block별로 나온 답 후보를 모두 활용해 pass@k나 best-block 관점에서 성능을 볼 수 있다. 논문 결과에서 any-block RiM은 특히 pass@8이 강하게 나타나며, 이는 memory block이 동일한 답을 반복하는 슬롯에 그치지 않고 서로 다른 refinement 지점을 제공한다는 해석을 뒷받침한다.
Stage 2의 손실 역시 표준 next-token prediction으로 표현된다. 차이는 target이 written reasoning step에서 final answer로 바뀐다는 점이다. 이 단순함은 RiM의 장점이자 한계다. 장점은 별도 verifier나 reward model 없이도 작은 모델에서 실행 가능한 학습 루프를 만든다는 점이다. 한계는 복잡한 과제에서 final answer supervision만으로 block별 내부 역할이 얼마나 안정적으로 유지되는지 아직 넓게 검증되지 않았다는 점이다.
| Table 2: RiM 구성 요소와 기능적 의미 | ||
|---|---|---|
| 구성 요소 | 논문에서의 역할 | 리뷰 관점의 해석 |
| Memory block | 고정 special token 묶음 | 중간 reasoning trace를 생성하지 않고 내부 상태를 저장하는 슬롯 |
| Stage 1 readout | 다음 reasoning step 예측 | 비어 있는 슬롯을 계산 가능한 작업공간으로 접지 |
| Stage 2 readout | 최종 답 예측 | block이 늘어날수록 답을 정제하는 훈련 목표 |
| Attention mask | target step 간 leakage 차단 | 병렬 supervision과 순차 정보 제약을 함께 만족 |
| Special token embedding update | 기존 vocabulary embedding freeze | pretrained 언어 지식을 흔들지 않고 새 workspace만 학습 |
Table 2에서 보듯 RiM의 설계는 새 아키텍처 블록을 크게 붙이는 방식보다 입력과 손실의 재배치에 가깝다. 이 점은 실무 적용에서 중요하다. 기존 LLM을 완전히 새로 설계하지 않아도 special token, attention mask, curriculum을 통해 latent workspace를 실험할 수 있기 때문이다. 다만 실제 서비스 모델에 적용하려면 prompt format, tokenizer 확장, supervised reasoning trace 확보, block 수 예산을 모두 운영 파라미터로 관리해야 한다.
4. 실험 설정: GSM8K-Aug, GSM8K, GSM-Hard와 비교 기준
4.1 데이터셋과 모델 규모
논문은 RiM을 넓은 도메인으로 즉시 확장하기보다, latent reasoning 선행 연구가 많이 사용한 수학 추론 환경에서 먼저 검증한다. 훈련 데이터는 GSM8K-Aug이며, 약 38만 5천 개의 grade-school math question과 reasoning step을 포함한다. 평가 데이터는 in-distribution 성격의 GSM8K와, 더 어려운 변형 문제로 구성된 out-of-distribution 성격의 GSM-Hard다.
모델은 GPT-2와 Llama-3.2 계열을 사용한다. 이 선택은 두 가지 의미를 갖는다. 첫째, latent reasoning이 특정 최신 대형 모델 하나에만 맞는 현상인지 분리해 볼 수 있다. 둘째, 작은 모델에서 memory block이 실제로 작동하는지 확인할 수 있다. 작은 모델은 reasoning token을 길게 생성할 때 비용 대비 성능이 더 민감하게 흔들리므로, RiM의 latency trade-off를 보기 좋다.
| Table 3: GSM8K-Aug reasoning step 길이 분포 | ||
|---|---|---|
| Reasoning steps | Training samples | Share |
| 1 | 62,908 | 16% |
| 2 | 143,578 | 37% |
| 3 | 104,249 | 27% |
| 4 | 48,198 | 13% |
| 5 | 17,906 | 5% |
| 6 | 5,666 | 1% |
| 7 | 2,359 | 1% |
| 8 이상 | 755 | 0%대 |
| Total | 385,620 | 100% |
Table 3은 RiM의 Stage 1 supervision이 어떤 분포 위에서 작동하는지 보여준다. 대부분의 훈련 샘플은 2~4개의 reasoning step에 몰려 있고, 8개 이상 step은 매우 적다. 따라서 RiM이 배운 working memory는 짧고 중간 길이의 수학 풀이에는 강하게 접지되어 있지만, 매우 긴 도구 사용 trajectory나 수십 단계 탐색으로 바로 일반화된다고 말하기는 어렵다. 이 점은 7절과 8절에서 한계로 다시 중요해진다.
4.2 비교 기준과 checkpoint 선택
비교 기준은 SFT w/o CoT, SFT w/ CoT, Coconut, DART 등으로 구성된다. SFT w/o CoT는 질문에서 직접 답을 내는 가장 빠른 baseline이고, SFT w/ CoT는 명시적 reasoning trace를 생성하는 강한 text reasoning baseline이다. Coconut은 continuous thought를 autoregressive로 생성하는 latent baseline이므로 RiM과 가장 직접적으로 비교된다. DART는 filler token을 reasoning에 쓰기 위한 별도 self-distillation 계열 방법으로, special token computation이라는 관점에서 관련이 깊다.
논문은 평가 프로토콜에서도 한 가지 중요한 선택을 한다. 선행 연구 일부는 여러 checkpoint를 평가한 뒤 test set에서 가장 높은 값을 보고하는 식으로 selection overfitting 위험을 만든다. 저자들은 이를 피하기 위해 16 split cross-validation checkpoint selection을 사용한다. 각 split에서 GSM8K 일부를 selection용으로 떼어 두고, 그 기준으로 checkpoint를 선택한 뒤 나머지 평가를 수행한다. 이 방식은 수치가 조금 덜 화려해 보일 수 있지만, 방법 간 비교의 신뢰도를 높인다.
4.3 latency 측정의 의미
RiM 논문에서 accuracy만큼 중요한 축은 time-to-first-token(TTFT)이다. CoT는 최종 답 앞에 많은 token을 생성하므로 첫 답을 보기까지 시간이 길어진다. Coconut은 자연어 reasoning token을 줄이지만 continuous thought를 순차적으로 만들어야 하므로 latency가 남는다. RiM은 memory block이 입력 안에 고정되어 있고 forward pass에서 병렬 처리되므로, 답변 token 생성 시작 시간은 직접 답변 SFT에 가까워진다.
이 측정은 실제 서비스 환경에서 특히 중요하다. 사용자는 모델 내부가 몇 단계 추론했는지보다 응답이 언제 시작되는지, 그리고 추가 비용이 어느 정도인지에 민감하다. RiM이 제안하는 방향은 reasoning capacity를 늘리면서도 사용자에게 보이는 token stream을 길게 만들지 않는 것이다. 따라서 이 논문은 reasoning 성능 연구인 동시에, 추론 serving 지연 시간을 줄이는 inference design 연구로도 읽힌다.
5. 주요 실험 결과: 성능과 latency trade-off
5.1 final-block RiM의 핵심 성능
Table 4는 final-block RiM이 직접 답변 SFT와 Coconut을 상대로 어느 위치에 서는지 보여준다. GPT-2에서는 SFT w/o CoT가 GSM8K greedy 15.4%에 머무는 반면, Coconut w/ Stage 0은 31.1%, RiM final block은 33.6%를 기록한다. GSM-Hard에서도 RiM은 7.8%로 Coconut의 7.1%보다 높다. Llama-3.2-1B에서는 SFT w/o CoT가 GSM8K 23.9%, Coconut이 36.9%, RiM이 42.1%로 더 큰 차이를 낸다.
| Table 4: 주요 결과 요약. RiM final block은 직접 답변 latency에 가깝게 유지되면서 Coconut 대비 정확도가 높다. | |||||||
|---|---|---|---|---|---|---|---|
| Model | Method | Variant | TTFT(ms) | GSM8K Greedy | GSM8K Pass@8 | GSM-Hard Greedy | GSM-Hard Pass@8 |
| GPT-2 | SFT | w/o CoT | 7.6 | 15.4±0.2 | 33.3±0.3 | 3.5±0.1 | 7.6±0.1 |
| GPT-2 | Coconut | w/ Stage 0 | 53.4 | 31.1±0.2 | 45.0±0.2 | 7.1±0.0 | 10.7±0.1 |
| GPT-2 | RiM | Final block | 7.6 | 33.6±0.2 | 49.1±0.2 | 7.8±0.1 | 11.2±0.1 |
| Llama-3.2-1B | SFT | w/o CoT | 16.1 | 23.9±0.2 | 41.7±0.3 | 5.3±0.1 | 9.5±0.1 |
| Llama-3.2-1B | Coconut | w/ Stage 0 | 108.3 | 36.9±0.2 | 51.1±0.2 | 8.5±0.0 | 12.2±0.0 |
| Llama-3.2-1B | RiM | Final block | 16.1 | 42.1±0.2 | 56.1±0.3 | 10.5±0.0 | 13.8±0.0 |
Table 4의 핵심은 RiM이 정확도를 높이는 동시에 TTFT를 늘리지 않는다는 점이다. GPT-2에서 Coconut은 53.4ms, RiM은 7.6ms로 표시된다. Llama-3.2-1B에서도 Coconut은 108.3ms, RiM은 16.1ms다. 물론 이 표의 TTFT는 특정 실험 환경과 입력 형식에 묶인 값이지만, 논문이 주장하는 구조적 차이는 분명하다. RiM은 reasoning state를 생성 token stream으로 만들지 않아 첫 답 token까지의 경로를 짧게 유지한다.
5.2 any-block readout과 pass@8의 차이
final-block만 보면 RiM은 빠른 latent workspace로 읽힌다. 하지만 any-block readout을 보면 또 다른 성격이 드러난다. 각 memory block 뒤의 답을 후보로 읽으면, 같은 forward pass 안에서 여러 refinement 지점이 생긴다. GPT-2에서 SFT w/ CoT는 GSM8K pass@8 57.0%를 기록하지만, RiM any-block은 78.1%까지 올라간다. Llama-3.2-1B에서도 RiM any-block은 pass@8 76.8%로 SFT w/ CoT 64.7%보다 높다.
| Table 5: any-block readout 결과. 작은 모델에서는 pass@8 이득이 두드러지지만 3B에서는 SFT w CoT가 더 강한 축도 남는다. | |||||||
|---|---|---|---|---|---|---|---|
| Model | Method | Variant | TTFT(ms) | GSM8K Greedy | GSM8K Pass@8 | GSM-Hard Greedy | GSM-Hard Pass@8 |
| GPT-2 | SFT | w/ CoT | 213.7 | 39.8±0.2 | 57.0±0.2 | 8.4±0.0 | 12.9±0.1 |
| GPT-2 | RiM | Any block | 7.6 | 39.5±0.3 | 78.1±0.2 | 9.4±0.0 | 19.0±0.1 |
| Llama-3.2-1B | SFT | w/ CoT | 420.3 | 49.1±0.4 | 64.7±0.3 | 11.2±0.1 | 15.3±0.1 |
| Llama-3.2-1B | RiM | Any block | 16.1 | 51.4±0.2 | 76.8±0.1 | 13.0±0.0 | 19.6±0.1 |
| Llama-3.2-3B | SFT | w/ CoT | 754.4 | 66.9±0.2 | 78.3±0.4 | 19.0±0.1 | 24.4±0.3 |
| Llama-3.2-3B | RiM | Any block | 27.9 | 57.3±0.1 | 71.8±0.2 | 13.8±0.1 | 18.2±0.0 |
Table 5는 RiM을 무조건 CoT보다 낫다고 읽지 않게 만든다. GPT-2와 Llama-3.2-1B에서는 any-block RiM이 매우 강하지만, Llama-3.2-3B의 greedy와 pass@8에서는 SFT w/ CoT가 더 높은 수치를 보인다. 이는 RiM이 장문 CoT의 모든 장점을 대체한다기보다, 작은 모델과 낮은 latency 예산에서 특히 매력적인 대안을 제공한다는 뜻이다. 고비용 reasoning을 허용할 수 있는 상황에서는 explicit trace가 여전히 강력할 수 있다.
Figure 4: RiM과 SFT 비교: accuracy-latency 평면에서 RiM이 직접 답변 속도에 가까운 위치를 차지한다.
Figure 5는 RiM과 SFT 계열의 정확도·latency 관계를 시각적으로 비교한다. 논문은 RiM이 CoT 없이 직접 답을 내는 SFT보다 높은 정확도를 제공하면서도, 긴 reasoning trace를 생성하는 SFT w CoT의 latency를 피한다고 보고한다. 이 그림은 RiM의 주장이 accuracy만의 우위가 아니라, 같은 응답 지연 예산에서 더 나은 reasoning capacity를 얻는 trade-off라는 점을 보여 준다.
Figure 5: RiM과 Coconut 비교: latent reasoning baseline 대비 순차 thought generation 비용을 줄이는 방향을 보여 준다.
Figure 6은 Coconut과 RiM의 비교를 통해 latent reasoning 내부에서도 실행 방식이 다르다는 점을 드러낸다. Coconut은 자연어 token을 피하더라도 continuous thought를 단계별로 생성해 다음 단계에 되먹인다. RiM은 fixed memory block을 입력 안에 두고 병렬 forward pass로 처리하므로, latent reasoning을 유지하면서 autoregressive 중간 표현 생성의 지연을 줄이는 쪽으로 이동한다.
5.3 latency 표가 말하는 운영적 의미
| Table 6: Llama-3.2-1B에서 GSM8K 질문당 inference latency | |||
|---|---|---|---|
| Method | Generated tokens | Average wall-clock time | 추가 지연 |
| SFT w/o CoT | 3.1 | 126.0±0.5ms | 기준 |
| SFT w/ CoT | 36.7 | 1108.7±3.0ms | +982.7ms |
| Coconut | 3.1 | 304.7±0.9ms | +178.7ms |
| RiM | 3.1 | 126.5±0.5ms | +0.5ms |
Table 6은 블로그 독자가 가장 바로 이해할 수 있는 수치다. SFT w/ CoT는 평균 생성 token이 36.7개로 늘고 wall-clock time도 1108.7ms로 커진다. Coconut은 출력 token 수가 3.1개로 줄어도 내부 continuous thought 처리 때문에 304.7ms가 걸린다. RiM은 126.5ms로 SFT w/o CoT의 126.0ms와 사실상 같은 수준이다. 논문이 말하는 “fixed memory block은 한 번에 처리된다”는 설명이 실제 latency 표에서 강하게 드러난다.
물론 이 숫자를 모든 GPU, batch size, serving engine에 그대로 옮길 수는 없다. 하지만 상대 비교는 유효한 신호다. CoT 방식은 token stream이 길어지는 만큼 사용자 지연이 커지고, Coconut 방식은 token stream은 짧아도 latent step이 순차적이면 내부 지연이 남는다. RiM은 중간 reasoning capacity를 입력 쪽 fixed block으로 이동시켜, 사용자가 보는 답변 시작 시간을 직접 답변 모델과 가깝게 유지한다.
6. 추가 분석 및 Ablation Study: memory block 표현, 예산 변화, attention mask
6.1 memory block representation은 실제로 달라지는가
RiM의 주장은 memory block이 단순 placeholder를 넘어 input-dependent latent workspace가 된다는 것이다. 이를 확인하기 위해 논문은 Llama-3.2-1B 모델의 memory block representation을 훈련 checkpoint별로 수집하고 PCA basis에 투영한다. 만약 block이 무시된다면 representation은 token identity와 위치에만 묶여 거의 같은 점에 머물 것이다. 반대로 문제별로 다른 내부 계산을 담는다면, 훈련이 진행되며 block별 trajectory가 벌어져야 한다.
Figure 6: memory block representation의 PCA trajectory: 훈련 전 collapsed 상태에서 질문별 구조를 가진 latent workspace로 퍼진다.
Figure 4는 RiM 논문에서 가장 중요한 진단 그림 중 하나다. 위쪽 행은 훈련 checkpoint가 진행되면서 memory block representation이 이동하는 궤적을 보여 주고, 아래쪽 행은 base model과 최종 모델의 표현 분포를 비교한다. 최종 모델에서는 질문별 representation이 더 넓게 퍼지며, 이는 memory block이 고정 token 위치를 넘어 입력별 작업 상태를 담기 시작했음을 시사한다. 즉 성능 향상이 단순 token 추가 효과만은 아니라는 근거가 된다.
저자들은 PCA 투영만으로 끝내지 않고, 멀리 떨어진 두 sample의 원래 representation 공간 cosine similarity도 확인한다. 이는 시각화상 분리가 단순 projection artifact에 머물지 않고 실제 hidden state 차이와 연결됨을 보이려는 장치다. 이런 분석은 RiM을 검은상자 성능 개선으로만 보지 않게 한다. memory block이 훈련 중 어떻게 조직되는지, 그리고 질문별로 다른 latent state를 형성하는지를 직접 추적하려는 시도다.
6.2 probe 기반 answer selection
| Table 7: memory block별 선형 probe separability와 answer selection | ||||||
|---|---|---|---|---|---|---|
| Metric | Memory block 1 | Memory block 2 | Memory block 4 | Memory block 6 | Memory block 8 | Probe-based selection |
| AUROC | 84.8±0.1 | 85.0±0.1 | 84.2±0.1 | 83.6±0.1 | 84.5±0.1 | 86.0±0.1 |
| AUPRC | 80.7±0.2 | 82.3±0.2 | 82.0±0.2 | 81.6±0.2 | 81.9±0.2 | 83.3±0.2 |
| Accuracy | - | - | - | - | - | 90.0±0.2 |
Table 7은 block별 representation이 답의 품질과 연결되는지를 probe로 본다. AUROC와 AUPRC가 80%대 중반을 유지한다는 것은 memory block representation 안에 답 선택과 관련된 신호가 들어 있음을 뜻한다. 특히 probe-based answer selection accuracy가 90.0%로 보고된 점은 any-block 후보를 그냥 모두 동일하게 취급하기보다, 어떤 block의 답을 채택할지 별도 selector를 붙일 수 있음을 보여 준다. 이 지점은 실제 deployment에서 중요하다.
운영 관점에서 probe는 latency를 크게 늘리지 않는 후처리로 붙일 수 있다. block별 답 후보가 있을 때 confidence, representation norm, probe score, answer transition 패턴을 함께 보아 최종 답을 고르는 식이다. 논문은 이 가능성을 완전히 제품화하지는 않지만, RiM이 단일 final answer generator보다 더 풍부한 내부 상태를 남긴다는 단서를 제공한다.
6.3 inference-time memory budget 변화
Figure 7: inference-time memory budget 변화에 따른 accuracy: memory block 수를 바꿔도 성능이 비교적 안정적으로 유지되는지 본다.
Figure 7은 inference 때 사용할 memory block 예산을 바꿨을 때 정확도가 어떻게 움직이는지 보여 준다. RiM이 특정 block 수에만 과적합됐다면 예산이 달라질 때 성능이 크게 무너질 수 있다. 논문은 block 수 변화에도 성능이 비교적 안정적으로 유지된다고 보고하며, 이는 memory block을 serving budget에 맞춰 조절할 수 있는 가능성을 만든다. 요청 난이도에 따라 block 수를 다르게 쓰는 adaptive serving의 출발점이다.
Figure 8: answer transition across memory blocks: block이 늘어날 때 답 후보가 어떻게 바뀌는지 추적한다.
Figure 8은 memory block이 단순 반복 슬롯인지, 답을 실제로 정제하는 단계인지 확인하는 그림이다. block을 지나며 answer가 바뀌는 패턴을 보면 모델이 같은 final answer를 복사만 하는지, 앞선 latent state를 바탕으로 후보를 수정하는지 판단할 수 있다. 이 분석은 any-block readout의 pass@k 이득이 어디서 나오는지 해석하는 데 필요하다. 또한 실패 후보가 어느 block에서 생기는지 추적하는 단서도 된다.
memory budget 분석은 RiM의 실용성을 평가할 때 핵심이다. 실제 시스템에서는 모든 요청에 같은 reasoning 예산을 줄 수 없다. 쉬운 질문에는 적은 block만 쓰고, 어려운 질문에는 더 많은 block을 허용하는 식의 adaptive budget이 필요하다. 논문은 아직 adaptive controller까지 제안하지는 않지만, block별 answer transition과 probe score를 결합하면 그 방향으로 확장할 수 있는 기초 신호를 제공한다.
6.4 bidirectional memory block attention ablation
Figure 9: bidirectional memory block attention 변형: 같은 block 내부 token끼리는 양방향으로 보고 block 사이 causal 구조는 유지한다.
Figure 9는 memory block 내부 attention을 어떻게 둘지에 대한 ablation 방향을 보여 준다. 기본 RiM은 causal 구조를 유지하지만, 같은 block 내부 token끼리 bidirectional attention을 허용하면 block 안의 작은 작업공간이 더 밀도 있게 상호작용할 수 있다. 다만 block 사이 정보 흐름은 여전히 causal하게 유지되므로, 전체 reasoning progression이 뒤섞이는 문제를 피하려는 설계다.
이 ablation은 memory block을 “토큰 몇 개 추가”로만 보지 않게 한다. block 내부 token이 서로 어떻게 정보를 교환하는지, block 간에는 어떤 시간 방향을 유지할지에 따라 latent workspace의 성격이 달라질 수 있다. 특히 긴 과제에서는 block 내부 bidirectional mixing이 더 강한 작업공간을 만들 수 있지만, 잘못 설계하면 답을 조기 고정하거나 block별 다양성을 줄일 위험도 있다.
3.4 Algorithm 관점에서 RiM 실행 흐름 풀어 읽기
RiM을 실제 학습 루프로 풀어 보면 첫 단계는 supervised trace를 block 단위 target으로 재배열하는 일이다. 원래 CoT 데이터에는 질문, 여러 reasoning step, 최종 답이 있다. RiM Stage 1은 이 trace를 그대로 출력하게 하지 않고, 질문 뒤에 $K$개의 memory block을 둔 다음 block $k$ 뒤 readout이 step $k+1$ 또는 그에 대응하는 reasoning segment를 예측하도록 만든다. 이 재배열은 데이터 자체를 바꾸기보다, 어느 위치의 hidden state가 어떤 target을 책임지는지를 바꾸는 작업이다. 따라서 구현자는 모델 구조보다 attention mask와 label mask를 더 조심해야 한다.
두 번째 단계는 추론 때 쓸 입력 형식을 고정하는 일이다. 질문이 들어오면 모델은 reasoning sentence를 생성하지 않고, 질문 token 뒤에 정해진 수의 memory block을 붙인 입력을 한 번 통과시킨다. final-block 방식에서는 마지막 block 뒤 readout에서 답을 읽는다. any-block 방식에서는 각 block 뒤 readout을 후보 답으로 보고, sampling 또는 selection rule로 답을 고른다. 이 차이는 사용자에게 보이는 출력 token 수를 늘리지 않으면서 내부 후보를 여러 개 만들 수 있다는 점에서, best-of-N sampling과 비슷한 목적을 더 낮은 latency로 달성하려는 시도다.
세 번째 단계는 budget을 정하는 일이다. block 수 $K$와 block 내부 memory token 수 $M$은 곧 내부 계산 예산이다. $K$를 늘리면 더 많은 refinement point가 생기지만, 입력 sequence 길이와 forward cost가 늘어난다. $M$을 늘리면 block 하나의 표현 용량이 커질 수 있지만, 논문은 기본 $M=2$로도 의미 있는 결과를 낸다. 이 선택은 RiM이 거대한 latent scratchpad를 붙이는 방법이라기보다, 아주 작은 fixed workspace를 강하게 supervised하면 어느 정도 계산 역할을 배울 수 있다는 주장에 가깝다.
네 번째 단계는 답을 읽는 정책이다. 논문은 final block, any block, best block, probe-based selection처럼 여러 readout 방식을 비교한다. 이들은 같은 학습된 memory representation을 서로 다른 운영 정책으로 사용하는 방법이다. final block은 가장 단순하고 deterministic한 serving에 적합하다. any block은 pass@k와 후보 다양성을 보기 좋다. probe-based selection은 block representation 자체가 어떤 답을 신뢰할 만한지 알려 줄 수 있는지를 묻는다. 따라서 RiM의 전체 알고리즘은 학습 방법과 readout 정책이 분리된 구조로 이해하는 편이 정확하다.
3.5 왜 fixed token이 계산 공간이 될 수 있는가
처음 보면 고정 special token을 붙이는 것만으로 작업 기억이 생긴다는 주장이 이상하게 보일 수 있다. 하지만 Transformer에서 token의 의미는 token id만으로 결정되지 않는다. 같은 token이라도 앞선 context, positional relation, attention pattern, layer별 transformation에 따라 contextual representation은 달라진다. RiM은 이 성질을 활용한다. <m> token의 표면형은 고정되어 있지만, 그 hidden state는 질문과 이전 block을 보고 만들어진다. 결국 memory block은 입력마다 새로 계산되는 latent variable처럼 행동할 수 있다.
차이는 훈련 신호다. 일반 filler token은 모델에게 “여기서 뭔가 계산해도 된다”는 빈칸만 준다. RiM Stage 1은 “이 빈칸을 지나면 다음 reasoning step을 맞혀야 한다”는 강한 제약을 건다. block이 다음 step을 맞히려면 앞선 숫자, 연산, 조건을 압축해 담아야 한다. 이때 block의 역할은 문장을 생성하는 것이 아니라, 다음 step을 생성할 수 있는 충분 통계량에 가까워진다. 논문이 PCA와 probe를 넣은 이유도 바로 이 hidden state가 실제로 문제별 정보를 담는지 보여 주기 위해서다.
이 설명은 모델 내부를 과도하게 의인화하지 않으면서도 “working memory”라는 말을 기능적으로 해석하게 해 준다. RiM의 memory block은 인간의 기억 구조와 동일하다는 뜻이 아니다. 입력 안에 존재하고, 질문별로 contextualized되며, 이후 readout이 중간 step이나 final answer를 예측하는 데 쓰이는 계산 슬롯이라는 뜻이다. 따라서 논문의 인지과학 비유는 구현 원리 자체라기보다, 외부 발화와 내부 조작을 분리하자는 설계 원칙으로 받아들이는 편이 안전하다.
4.4 평가 프로토콜에서 놓치기 쉬운 세부 지점
논문이 cross-validation checkpoint selection을 강조하는 이유는 latent reasoning 비교에서 생각보다 중요하다. CoT, Coconut, RiM처럼 훈련 안정성이 다른 방법을 비교할 때, test set에서 가장 좋은 checkpoint를 고르면 특정 방법에 유리한 우연을 성능처럼 보고할 수 있다. RiM은 16개 split에서 selection용 GSM8K subset을 따로 두고, 그 선택 기준을 평가와 분리한다. 이 절차는 숫자를 낮출 수 있지만, 재현 가능한 비교에는 더 안전하다.
또 하나의 세부 지점은 answer prefix 강제다. 논문은 SFT w/o CoT와 latent reasoning 방법에서 “The final answer is boxed” 형식의 prefix를 강제해 final-answer generation만 평가한다. 이 설정은 모델이 긴 설명을 새로 만들면서 추가 reasoning을 하게 되는 경로를 막는다. RiM의 목표가 내부 memory block을 쓰는 것이라면, 평가에서도 출력 쪽에서 몰래 추가 reasoning을 생성하지 않게 통제해야 한다. 이 작은 설정이 accuracy-latency 비교의 공정성에 영향을 준다.
GSM-Hard를 함께 쓰는 것도 중요하다. GSM8K는 훈련 데이터와 분포가 가까운 in-distribution 평가이고, GSM-Hard는 더 어려운 out-of-distribution 변형이다. RiM이 단순히 GSM8K-Aug step pattern을 외워서 답을 맞히는 것인지, 난이도가 올라간 수학 문제에서도 latent workspace가 도움이 되는지 보려면 둘을 나눠야 한다. Table 4와 Table 5에서 RiM은 GSM-Hard에서도 일부 이득을 보이지만, 절대 수치는 여전히 낮다. 이는 latent workspace가 강력한 신호를 주더라도 작은 모델의 OOD reasoning 한계가 남는다는 뜻이다.
평가 지표도 하나로 충분하지 않다. Greedy accuracy는 가장 직접적인 답 품질을 보여 주지만, any-block readout의 가치는 pass@8에서 더 잘 나타난다. TTFT는 사용자 지연을 보여 주고, wall-clock latency는 내부 순차 계산 비용을 보여 준다. probe AUROC와 AUPRC는 hidden state가 답 선택 정보를 담는지 보여 준다. RiM을 평가하려면 이 지표들이 서로 같은 방향을 가리키는지 봐야 한다. 이 논문은 그런 점에서 단순 벤치마크 점수표보다 진단 항목을 비교적 잘 배치했다.
5.4 결과를 모델 크기별로 다시 읽기
GPT-2 결과는 RiM의 minimum viable signal을 보여 준다. 아주 강한 최신 LLM이 아니어도 memory block에 reasoning step supervision을 걸면 direct answer SFT보다 훨씬 나은 결과가 나온다. GPT-2 SFT w/o CoT의 GSM8K greedy 15.4%와 RiM final block 33.6%의 차이는 작지 않다. 이 차이는 fixed workspace가 단순히 고급 모델의 여분 capacity를 쓰는 장식이 아니라, 작은 모델에서도 계산 구조를 바꾸는 효과가 있음을 시사한다.
Llama-3.2-1B 결과는 RiM의 가장 설득력 있는 구간이다. SFT w/o CoT보다 accuracy가 높고, Coconut보다 TTFT가 훨씬 낮으며, SFT w/ CoT보다 latency가 압도적으로 작다. 이 모델 크기는 실제 edge나 저비용 serving 논의와도 연결된다. 1B급 모델에서 장문 CoT를 매번 생성하는 것은 비용 대비 부담이 크다. RiM은 이런 상황에서 내부 계산 슬롯을 조금 늘려 직접 답변의 속도에 가까운 추론을 시도하는 방법으로 읽힌다.
Llama-3.2-3B 결과는 더 조심스럽게 읽어야 한다. any-block RiM은 작은 모델에서 강했지만, 3B에서는 SFT w/ CoT가 일부 지표에서 더 좋은 값을 낸다. 이는 모델 규모가 커질수록 명시적 reasoning trace를 학습하고 생성하는 능력이 강해져, hidden workspace만으로는 모든 이득을 가져오기 어려울 수 있음을 보여 준다. 따라서 RiM은 “CoT의 대체재”라기보다, latency와 audit 요구가 서로 다른 상황에서 선택할 수 있는 추가 inference mode로 보는 것이 좋다.
이 모델 크기별 패턴은 후속 연구 설계에도 영향을 준다. 작은 모델에서는 fixed memory block 자체가 큰 이득을 줄 수 있고, 큰 모델에서는 CoT와 RiM의 hybrid가 더 나을 수 있다. 예를 들어 간단한 문제는 RiM으로 빠르게 처리하고, block별 답이 불안정하거나 verifier가 낮은 점수를 주면 explicit CoT로 전환하는 방식이다. 이는 latency 절감과 해석 가능성 사이의 균형을 운영 정책으로 옮기는 방향이다.
5.5 pass@8 이득을 후보 다양성 관점으로 보기
RiM any-block의 pass@8 이득은 단순히 여러 번 sampling해서 얻은 이득과 다르게 볼 필요가 있다. 일반 pass@k는 같은 모델을 여러 번 샘플링해 후보를 얻는다. RiM any-block은 하나의 입력 안에서 서로 다른 memory block readout을 후보로 읽는다. 이 차이는 후보들이 같은 generation distribution의 독립 샘플인지, 아니면 block depth에 따른 refinement 상태인지의 차이를 만든다. 논문은 answer transition 그림을 통해 block이 늘어날수록 답이 바뀌는 과정을 따로 본다.
이 관점은 VPO에서 다뤘던 test-time search post-training과도 연결된다. VPO는 후보 집합의 reward-space diversity를 훈련해 best@k search가 더 잘 작동하게 했다. RiM은 reward-vector diversity를 직접 훈련하지는 않지만, block별 readout을 통해 하나의 forward pass 안에서 여러 refinement 후보를 얻는다. 따라서 둘을 결합하면 흥미로운 실험이 가능하다. memory block마다 다른 reward trade-off나 reasoning strategy를 유도하고, search가 그 후보를 고르는 방식이다.
다만 pass@8이 높다고 해서 사용자에게 항상 8개 답을 보여 주면 된다는 뜻은 아니다. 실제 시스템은 최종적으로 하나의 답을 내야 하고, 틀린 후보를 포함한 candidate pool을 어떻게 다룰지 결정해야 한다. 이때 probe-based answer selection이나 verifier가 필요하다. RiM 논문은 linear probe 성능을 보여 주지만, production-grade selection policy까지 제시하지는 않는다. 따라서 pass@8 이득은 잠재력으로 읽되, 단일 답 품질과 후보 선택 비용을 함께 평가해야 한다.
6.5 memory block을 관찰 가능한 로그로 바꾸는 방법
RiM의 중간 상태는 hidden representation이므로 그대로는 사람이 읽을 수 없다. 하지만 완전히 관찰 불가능한 것은 아니다. block별 final answer, answer transition 여부, probe confidence, representation distance, entropy, verifier score를 남기면 latent workspace의 상태 변화를 간접적으로 로그화할 수 있다. 특히 Figure 8처럼 block별 답 변화가 있다면, 어느 block에서 답이 처음 안정화되는지 기록할 수 있다. 이는 explicit CoT 없이도 최소한의 reasoning trace receipt를 만드는 방법이 된다.
검색 에이전트나 코딩 에이전트에 RiM을 적용한다면 이 로그 설계가 더 중요해진다. 예를 들어 block 1은 문제 요약, block 2는 후보 파일, block 3은 수정 전략, block 4는 test 결과 반영처럼 명시적으로 이름 붙일 수는 없더라도, block별 output probe가 어떤 상태를 담는지 후처리 classifier로 추정할 수 있다. 이 classifier가 안정적으로 작동한다면 latent memory를 black-box로만 두지 않고, 운영 dashboard에서 “어느 단계까지 확신이 올라갔는지”를 볼 수 있다.
이 접근은 기존 wiki의 verification trace receipt 개념과도 닿는다. 모든 reasoning을 공개하지 않더라도, 검증 결과와 다시 열 첫 지점, 미확인 영역을 짧게 남기면 다음 실행이 안전해진다. RiM에서는 block별 confidence와 answer transition을 그런 receipt로 삼을 수 있다. 단, 이 정보가 실제 오류 원인과 얼마나 잘 대응하는지는 별도 benchmark가 필요하다. latent representation 로그가 그럴듯해 보여도 실패 설명을 보장하지는 않는다.
6.6 수식과 objective를 구현 관점으로 다시 정리
논문의 수식은 크게 두 손실로 요약된다. Stage 1은 $\mathcal{L}_{S1}$로, memory block까지 본 뒤 다음 reasoning step을 맞히는 negative log-likelihood다. Stage 2는 $\mathcal{L}_{S2}$로, memory block prefix별 final answer likelihood를 합산한다. 두 손실 모두 표준 next-token prediction 형태를 유지한다. 이 점은 구현 난이도를 낮추지만, label alignment와 mask 구성 오류가 성능에 크게 영향을 줄 수 있음을 뜻한다.
Stage 1의 label alignment에서는 reasoning step 개수와 memory block 개수가 맞지 않는 사례를 처리해야 한다. GSM8K-Aug에는 1개부터 13개 step까지 다양한 길이가 있으며, block 수 $K$가 고정되어 있다면 짧은 trace와 긴 trace를 어떻게 padding하거나 truncate할지 정해야 한다. Stage 2에서는 각 block 뒤 final answer를 예측하므로, answer target을 여러 위치에 반복 배치한다. 이 반복 배치가 너무 강하면 초기 block이 성급한 답을 내도록 과하게 압박할 수 있고, 너무 약하면 late block만 쓰게 될 수 있다.
attention mask는 RiM 구현에서 특히 민감하다. target reasoning step끼리 서로를 보면 Stage 1이 쉬운 leakage 문제가 된다. 반대로 memory block이 필요한 context를 충분히 보지 못하면 latent workspace가 형성되지 않는다. 논문은 memory block이 질문과 이전 block을 보도록 하며, target step은 이전 block과 선택적으로 질문을 보되 다른 target step은 보지 못하게 한다. 이 제약은 병렬 supervision과 정보 누수 방지를 동시에 만족시키려는 설계다.
따라서 RiM을 재현할 때는 모델 코드보다 데이터 collator와 mask builder를 먼저 검증해야 한다. 각 token 위치가 어떤 label을 갖는지, loss가 어디에 걸리는지, readout 위치가 block boundary와 맞는지, special token embedding만 업데이트되는지 확인해야 한다. 논문이 단순한 objective를 쓰는 만큼, 작은 구현 실수도 “memory block이 작동하지 않는다”는 잘못된 결론으로 이어질 수 있다.
7. 한계점 및 향후 연구 방향: 수학 벤치마크를 넘어선 검증 필요
RiM의 가장 큰 한계는 검증 범위가 수학 reasoning benchmark에 강하게 묶여 있다는 점이다. GSM8K-Aug는 reasoning step supervision을 제공하고, 답도 비교적 명확하다. 이 환경은 memory block을 접지하기 좋지만, 실제 agent task는 관찰, 도구 호출, 파일 상태, 웹 검색 결과, 장기 context가 얽힌다. 그런 과제에서는 “다음 reasoning step”이 텍스트 풀이 한 줄로 깔끔하게 정의되지 않을 수 있다.
두 번째 한계는 training signal의 비용이다. RiM은 추론 시 빠르지만, Stage 1을 위해 step-level reasoning supervision이 필요하다. 수학 문제에서는 풀이 단계가 데이터로 제공되거나 생성될 수 있지만, 코딩 에이전트나 연구 에이전트에서는 좋은 중간 단계 label을 얻기 어렵다. 잘못된 trace로 Stage 1을 학습하면 memory block이 오류를 압축해 저장할 수도 있다. 따라서 RiM을 더 넓은 도메인에 적용하려면 trace 품질 관리, synthetic trace 검증, 실패 step filtering이 함께 필요하다.
세 번째 한계는 해석 가능성이다. CoT는 위험한 hallucination을 만들 수 있지만, 적어도 사람이 읽을 수 있는 중간 문장이 있다. RiM은 중간 상태를 hidden representation으로 옮기므로, 성능이 좋아져도 왜 그런 답을 냈는지 직접 확인하기 어렵다. 논문은 PCA와 linear probe를 제공하지만, 이것은 전체 reasoning 과정을 설명하는 수준까지 가지는 않는다. 특히 안전이나 감사가 중요한 환경에서는 latent workspace가 black-box로 남는 비용을 따로 평가해야 한다.
향후 연구 방향으로는 세 가지가 자연스럽다. 첫째, Stage 2를 supervised final-answer loss만으로 두지 않고, final-answer reward나 verifier feedback을 결합하는 방법이다. 저자들도 RL with final-answer rewards 가능성을 언급한다. 둘째, memory block 수를 요청별로 조절하는 adaptive budget controller다. 쉬운 문제는 적은 block으로 끝내고, block별 answer transition이 불안정할 때만 더 많은 block이나 explicit CoT를 쓰는 방식이 가능하다. 셋째, code, tool-use, long-context search, multimodal reasoning에서 memory block이 어떤 state를 담는지 비교해야 한다.
마지막으로 RiM은 안전 측면에서 양면성이 있다. 중간 추론을 외부로 드러내지 않으면 사용자에게 불필요한 장문 reasoning을 노출하지 않아도 되고, latency도 낮아진다. 반대로 내부 작업공간이 보이지 않기 때문에 정책 위반 reasoning이나 잘못된 중간 가정이 어디서 생겼는지 감사하기 어려울 수 있다. 따라서 RiM류 latent reasoning은 deployment 단계에서 probe, trace summary, fallback CoT, verifier audit를 함께 설계해야 한다.
7.1 도메인 확장 시 가장 먼저 확인할 실험
RiM을 다음 단계로 검증한다면 가장 먼저 해야 할 일은 reasoning step이 명확한 수학 문제에서 벗어나, 상태 추적이 필요한 과제로 옮기는 것이다. 예를 들어 코드 디버깅에서는 실패 로그를 읽고, 의심 파일을 고르고, 수정한 뒤 테스트 결과를 반영해야 한다. 이 과정의 중간 상태는 수식 풀이 한 줄처럼 정리되지 않는다. memory block이 이런 상태를 담는다면 block별 probe는 파일 후보, 실패 유형, 수정 방향 같은 정보를 예측할 수 있어야 한다. 이런 probe가 없다면 RiM의 working memory 주장은 수학 trace의 압축에 머물 가능성이 크다.
검색 에이전트도 좋은 시험장이다. LongSeeker식 장기 검색에서는 이미 읽은 문서, 버린 branch, 남은 질문, 근거 snippet을 계속 관리해야 한다. RiM을 여기에 붙인다면 memory block이 검색 observation을 압축하는지, 혹은 단순히 마지막 문서 제목만 반영하는지 확인해야 한다. 평가도 final answer accuracy 하나로 끝나면 부족하다. evidence retention, deleted branch safety, citation consistency, rollback 가능성까지 함께 봐야 한다. 이런 지표가 추가될 때 RiM은 추론 benchmark를 넘어 agent memory 연구와 실제로 만날 수 있다.
멀티모달 과제에서는 또 다른 문제가 생긴다. MemEye에서 다뤘던 것처럼 agent memory는 텍스트 caption만으로 대체되지 않는 원본 시각 증거를 다뤄야 한다. RiM memory block이 이미지 region, 시간 상태, UI state를 hidden representation으로 담는다면, 그 representation이 어떤 시각 증거를 보존하는지 별도 probe가 필요하다. 특히 시각 정보가 답에 결정적일 때, latent workspace가 caption shortcut을 학습하는지 원본 visual feature를 유지하는지 구분해야 한다.
7.2 safety와 audit 관점의 숙제
RiM의 장점은 중간 reasoning을 사용자에게 길게 노출하지 않는 데 있다. 하지만 안전 관점에서는 이것이 동시에 숙제다. 모델이 어떤 잘못된 가정으로 답을 만들었는지, policy violation에 가까운 내부 경로를 거쳤는지, 특정 shortcut을 사용했는지 사람이 바로 읽기 어렵다. CoT를 항상 공개하는 방식도 정답은 아니지만, 완전히 숨겨진 latent workspace만 두는 것도 감사 가능성을 낮춘다. 따라서 RiM류 방법은 final answer 옆에 최소한의 diagnostic receipt를 남기는 운영 설계를 함께 가져가야 한다.
한 가지 방법은 selective externalization이다. 기본 응답은 RiM으로 빠르게 만들고, block별 confidence가 낮거나 answer transition이 여러 번 흔들리거나 verifier가 불일치를 감지하면 그때만 짧은 reasoning summary를 생성한다. 이 summary는 전체 CoT를 그대로 노출하는 것이 아니라, 어떤 근거를 확인했고 어떤 가정이 남았는지 요약하는 형태가 될 수 있다. 이렇게 하면 latency와 audit 사이에서 요청별로 다른 정책을 적용할 수 있다.
또 다른 방법은 block representation 자체의 drift를 감시하는 것이다. 운영 중 입력 분포가 바뀌면 memory block representation이 훈련 때와 다른 영역으로 이동할 수 있다. 이때 final answer가 그럴듯해도 내부 workspace가 안정적이라고 볼 수 없다. PCA나 linear probe는 연구용 진단이지만, production에서는 embedding distance, confidence calibration, verifier disagreement를 함께 기록해 drift alert를 만들 수 있다. 이런 계층이 있어야 latent reasoning이 보이지 않는 비용을 줄일 수 있다.
7.3 학습 데이터 품질과 trace provenance
RiM Stage 1은 reasoning step supervision에 크게 의존하므로, trace provenance가 중요하다. 사람이 작성한 풀이, 모델이 생성한 synthetic trace, verifier로 필터링한 trace는 품질과 오류 패턴이 다르다. GSM8K-Aug처럼 정리된 데이터에서는 이 문제가 덜 보이지만, 실제 업무형 과제에서는 중간 step label이 부정확하거나 과도하게 장황할 수 있다. 그런 trace로 memory block을 접지하면, 모델은 올바른 계산 구조보다 데이터 생성기의 말버릇이나 shortcut을 latent workspace에 새길 수 있다.
따라서 후속 연구는 Stage 1 trace를 단순 input으로 취급하지 말고, trace 자체의 품질을 실험 변수로 분리해야 한다. 짧은 trace와 긴 trace, human trace와 model trace, verifier-passed trace와 raw trace를 나눠 memory block representation이 어떻게 달라지는지 봐야 한다. 또한 Stage 2에서 final answer loss만 넣었을 때 Stage 1에서 배운 구조가 얼마나 유지되는지도 확인해야 한다. 이 실험은 RiM이 trace distillation에 가까운지, 실제로 일반화 가능한 작업 기억을 얻는지 가르는 기준이 된다.
이 지점은 블로그 운영 관점에서도 중요하다. RiM을 읽고 바로 “CoT를 없애도 된다”고 결론 내리기보다, 어떤 supervised trace가 memory block을 만들었는지 먼저 봐야 한다. 좋은 internal workspace는 좋은 중간 신호에서 나온다. 중간 신호가 부실하면 숨겨진 workspace는 더 빠른 오류 압축기가 될 수 있다. 논문은 이 위험을 직접 크게 다루지는 않지만, Stage 1 구조가 강한 만큼 trace 품질 관리는 자연스러운 다음 질문으로 남는다.
7.4 serving 시스템에 붙일 때의 체크리스트
실제 serving에 RiM을 붙인다면 tokenizer 확장, prompt template, KV cache, batching 정책을 함께 봐야 한다. special token이 추가되면 tokenizer와 embedding table이 바뀌고, memory block 수가 요청별로 달라지면 batch 안 sequence length도 달라진다. TTFT가 직접 답변과 비슷하다는 논문 결과는 특정 실험 환경에서 얻은 값이므로, vLLM 같은 serving stack에서는 prefill cost, cache reuse, speculative decoding과의 상호작용을 따로 측정해야 한다.
또한 사용자별 또는 과제별 budget 정책이 필요하다. 모든 요청에 $K=8$ block을 붙일지, 쉬운 요청은 $K=2$로 시작할지, verifier 실패 시 $K$를 늘릴지 정해야 한다. 이 정책은 가격표와도 연결된다. RiM은 output token을 줄여 비용을 낮추지만, input-side memory block과 prefill compute는 여전히 비용을 만든다. 따라서 “출력 토큰이 짧다”는 사실만으로 전체 비용 절감을 단정하지 말고, prefill과 decoding을 나눠 측정해야 한다.
마지막으로 fallback UX가 필요하다. RiM이 확신하지 못하는 요청에서 모델이 바로 짧은 답만 내면 사용자는 오류 원인을 알기 어렵다. 이때 “추론 근거를 더 자세히 계산 중” 같은 긴 진행 문구보다, 근거 재검토나 explicit reasoning mode 전환을 조용히 수행하고 최종 답에 필요한 설명만 붙이는 편이 낫다. RiM은 내부 작업공간을 제공하지만, 사용자 경험은 여전히 명확한 답, 필요한 근거, 오류 시 재시도 경로로 설계되어야 한다.
RiM의 결과를 장기적으로 보려면 pretraining-time memory와 inference-time workspace를 구분해 추적하는 것도 필요하다. Memory Grafting은 외부 hidden-state bank를 통해 모델 용량을 늘리는 방향이고, RiM은 한 요청 안에서 중간 상태를 조작하는 방향이다. 두 흐름이 결합되면 모델은 장기 지식은 외부 latent bank에서 가져오고, 당장의 문제 풀이는 fixed memory block에서 처리할 수 있다. 이 조합은 매력적이지만, 지식 검색 오류와 추론 오류가 같은 hidden state 공간에서 섞일 위험도 함께 만든다.
그래서 나는 RiM을 읽을 때 “CoT를 대체하는 새 정답”보다, reasoning surface를 여러 층으로 나누는 설계로 보는 편이 더 생산적이라고 생각한다. 사용자가 보는 답변 surface, verifier가 읽는 trace surface, 모델 내부가 쓰는 memory block surface가 서로 다른 속도로 움직일 수 있다. 논문은 그중 내부 surface를 작고 빠르게 만드는 데 성공한 사례다. 다음 단계는 이 내부 surface가 언제 충분하고, 언제 외부 trace나 tool evidence로 올라와야 하는지 결정하는 정책을 만드는 일이다.
8. 내 해석: 약점 1 + 후속 제안 1
나는 이 논문을 reasoning을 더 길게 말하게 하는 흐름에 대한 실행 구조 차원의 반론으로 읽었다. Path-Lock Expert가 reasoning mode가 출력으로 새는 문제를 다뤘다면, RiM은 더 앞에서 “중간 계산을 꼭 출력 sequence에 올려야 하는가”를 묻는다. Memory Grafting이 외부 latent memory bank를 사전학습에 붙이는 방식이었다면, RiM은 같은 forward pass 안에서 고정 memory block을 질문별 내부 작업공간으로 바꾸는 방식이다. 이 연결이 좋다. 기존 wiki에서 다뤘던 memory와 reasoning 논문들이 대체로 검색, 출력 모드, 후학습 신호에 집중했다면, RiM은 입력 format과 attention mask만으로도 추론의 시간 구조를 바꿀 수 있음을 보여 준다.
다만 내가 걸리는 약점은 과제 범위다. GSM8K/GSM-Hard는 수학 reasoning 연구에서 표준적이고 비교가 쉽지만, RiM이 말하는 working memory는 더 넓은 작업 상태를 담을 때 진짜 가치가 커진다. 예를 들어 코딩 에이전트는 실패 로그, 수정한 파일, 테스트 결과, 사용자 제약을 기억해야 하고, 검색 에이전트는 이미 본 문서와 버린 branch를 구분해야 한다. 이런 과제에서는 reasoning step supervision이 짧은 수식 풀이처럼 깔끔하지 않다. 따라서 이 논문만으로 “LLM의 작업 기억을 열었다”고 말하기보다는, 수학 풀이 환경에서 fixed memory block이 latent computation을 배울 수 있음을 강하게 보인 결과로 제한해 읽는 편이 안전하다.
내가 이걸 확장한다면 먼저 hybrid externalization policy를 붙여볼 것 같다. 기본 경로는 RiM block으로 빠르게 답을 만들고, block별 probe confidence가 낮거나 answer transition이 계속 흔들릴 때만 explicit CoT 또는 tool trace를 외부화하는 방식이다. 이렇게 하면 RiM의 latency 장점을 살리면서도, 실패가 의심되는 경우에는 사람이 읽을 수 있는 근거를 확보할 수 있다. 특히 기존 LongSeeker식 working memory나 SAVeR식 belief verification과 결합하면, memory block representation, verifier score, external trace를 하나의 운영 로그로 묶어 latent reasoning의 감사 가능성을 높일 수 있다.
9. 결론: 내부 작업공간으로 옮겨 간 latent reasoning
RiM 논문은 중간 추론을 더 많이 생성하는 방향과 다른 길을 제시한다. 핵심은 reasoning step을 자연어로 길게 쓰거나 continuous thought로 순차 생성하는 대신, fixed memory block의 contextual representation을 내부 작업공간으로 쓰는 것이다. Stage 1은 그 block을 다음 reasoning step 예측으로 접지하고, Stage 2는 final answer refinement로 목표를 바꾼다. 이 두 단계는 표준 next-token prediction 위에 얹히므로, 방법론 자체는 복잡한 새 optimizer보다 데이터 구성과 mask 설계에 가깝다.
실험 결과는 논문 주장을 꽤 명확하게 뒷받침한다. RiM final block은 GPT-2와 Llama-3.2-1B에서 Coconut보다 높은 정확도를 보이면서 TTFT는 직접 답변 SFT와 같은 수준에 머문다. any-block readout은 작은 모델에서 pass@8을 크게 끌어올리고, representation 분석은 memory block이 질문별로 달라지는 latent state를 형성한다는 증거를 제공한다. latency 표에서 Llama-3.2-1B RiM이 126.5ms로 SFT w/o CoT 126.0ms와 거의 같은 값이라는 점은 이 논문의 실용적 메시지를 압축한다.
그럼에도 RiM을 범용 해결책으로 받아들이기에는 아직 이르다. reasoning step label이 있는 수학 과제에서의 성공과, 도구 사용·코드 수정·검색·멀티모달 관찰이 섞인 에이전트 과제에서의 성공은 다르다. 중간 reasoning이 hidden state로 들어가면 사용자에게 보이는 trace가 줄어드는 대신 감사 비용이 늘 수 있다. 따라서 RiM의 다음 검증은 더 어려운 benchmark에서 accuracy를 재는 것만으로 충분하지 않다. block별 confidence, answer transition, failure fallback, explicit trace 전환 정책까지 함께 평가해야 한다.
이 논문의 가치가 가장 잘 살아나는 지점은 응답 지연 예산이 빡빡한 reasoning이다. 장문 CoT를 허용하면 성능은 좋아질 수 있지만, 모든 요청에 그런 비용을 쓸 수는 없다. RiM은 내부 계산 슬롯을 입력 쪽으로 옮겨, 사용자에게 보이는 출력은 짧게 유지하면서도 모델 내부에는 더 넓은 계산 공간을 제공하려 한다. 이 방향은 앞으로 작은 모델, 온디바이스 모델, 실시간 에이전트 UI에서 특히 중요해질 가능성이 있다.
10. 요약 정리
- RiM은 chain-of-thought처럼 중간 풀이를 생성하지 않고, 고정된 special token memory block을 내부 latent workspace로 쓰는 방법이다.
- Stage 1은 각 memory block 뒤에서 다음 reasoning step을 예측하게 해 block representation을 작업 기억으로 접지한다.
- Stage 2는 step supervision을 제거하고 각 block 뒤에서 final answer를 예측하게 해 답을 점진적으로 정제한다.
- 논문은 GSM8K-Aug로 훈련하고 GSM8K/GSM-Hard에서 GPT-2와 Llama-3.2 계열 모델을 평가하며, SFT, CoT, Coconut, DART와 비교한다.
- final-block RiM은 GPT-2와 Llama-3.2-1B에서 Coconut보다 높은 정확도를 보이면서 TTFT는 직접 답변 SFT와 비슷한 수준으로 유지한다.
- any-block readout은 작은 모델에서 pass@8을 크게 높이며, block별 답 후보가 단순 반복을 넘어 refinement 지점을 제공한다는 해석을 가능하게 한다.
- PCA representation, linear probe, answer transition 분석은 memory block이 질문별 latent state를 형성한다는 간접 증거를 제공한다.
- 주요 한계는 수학 벤치마크와 step-level supervision에 의존한다는 점이며, tool-use agent나 long-context search로 확장하려면 별도 검증이 필요하다.
- 운영적으로는 RiM block confidence가 낮을 때만 explicit CoT나 verifier trace를 켜는 hybrid fallback 정책이 자연스러운 후속 방향이다.