대규모 언어 모델의 작업 기억 활용을 통한 잠재적 추론 능력 향상

Unlocking the Working Memory of Large Language Models for Latent Reasoning

대규모 언어 모델의 추론 능력을 향상시키기 위해, 일반적으로 최종 답변을 생성하기 전에 중간 토큰을 생성하여 연산량을 늘립니다. 그러나 이러한 방식은 추론 과정을 자동 회귀 방식으로 처리하게 되므로, 내부 연산과 외부 통신 간의 구분이 모호해집니다. 반면, 인간 인지 과정에서는 작업 기억을 사용하여 중간 생각을 외부로 드러내지 않고 정보를 내부적으로 저장하고 조작할 수 있습니다. 본 연구는 이러한 원리를 바탕으로, 추론 단계를 자동 회귀 방식으로 생성하는 대신 메모리 블록을 사용하는 '메모리 내 추론 (RiM)'이라는 새로운 잠재적 추론 방법을 제안합니다. 이러한 메모리 블록은 대규모 언어 모델의 작업 기억 용량을 활용할 수 있도록 설계된 특수 토큰들의 고정된 시퀀스입니다. 생성되는 것이 아닌 고정된 시퀀스이기 때문에, 단일 순전파 과정을 통해 처리될 수 있으며, 이를 통해 효율적인 잠재적 추론이 가능합니다. 이러한 메모리 블록을 효과적으로 활용하기 위해, 두 단계의 교육 과정(curriculum)을 사용했습니다. 먼저, 각 메모리 블록 이후에 명시적인 추론 단계를 예측하도록 학습시켜 메모리 블록을 '고정'시키고, 그 다음에는 단계별 지도 없이 각 메모리 블록 이후 최종 답변을 반복적으로 개선합니다. 다양한 종류와 크기의 대규모 언어 모델에 대한 실험 결과, RiM은 기존의 잠재적 추론 방법과 동등하거나 더 나은 성능을 보였으며, 동시에 자동 회귀 방식으로 생각을 생성하는 단점을 극복했습니다. 이러한 결과는 대규모 언어 모델이 작업 기억을 활용하여 효과적인 잠재적 추론 메커니즘을 구축할 수 있음을 보여줍니다.

To improve the reasoning capabilities of large language models, test-time compute is typically scaled by generating intermediate tokens before the final answer. However, this couples reasoning to autoregressive generation and thereby conflates internal computation with external communication. In contrast, human cognition can use working memory to hold and manipulate information internally without the need to externalize intermediate thoughts. Drawing on this principle, we introduce Reasoning in Memory (RiM), a latent reasoning method that replaces the autoregressive generation of reasoning steps with memory blocks. These memory blocks are fixed sequences of special tokens that unlock the working-memory capacity of large language models. Since they are fixed rather than generated, they can be processed in a single forward pass, enabling compute-efficient latent reasoning. To operationalize these memory blocks, we employ a two-stage curriculum. First, we ground them by predicting explicit reasoning steps after each memory block. Second, we discard this step-level supervision and iteratively refine the final answer after each memory block. Our experiments on reasoning benchmarks show that, across language models of different families and sizes, RiM matches or exceeds existing latent reasoning methods while avoiding the autoregressive generation of thoughts. These results demonstrate that large language models can be trained to use working memory as an effective mechanism for latent reasoning.