DiscoLoop: 다중 단계 추론을 위한 반복적인 이산 임베딩과 연속적인 은닉 상태

DiscoLoop: Looping Discrete Embeddings and Continuous Hidden States for Multi-hop Reasoning

대규모 언어 모델은 Chain-of-Thought (CoT) 방식으로 중간 단계를 외부화할 수 있을 때 많은 추론 작업에서 뛰어난 성능을 보입니다. 그러나 많은 질문에서는 모델이 답변을 생성하기 전에 여러 단계의 추론을 단일 순방향 과정 내에 통합해야 합니다. 본 연구는 대표적인 다중 단계 추론 과제인 2단계 추론을 통해 이 문제를 분석합니다. 표준 비순환 트랜스포머는 '깊이 지역 저장' 문제로 인해 초기 레이어에서 학습된 사실들이 두 번째 단계 정보 검색 시점에 사용 불가능하게 됩니다. Looped Transformer가 동일한 메모리를 재사용하여 이 문제를 완화하지만, 여전히 일반화 성능에 한계가 있습니다. 분석 결과, 남은 병목 현상은 표현 방식과 관련된 것으로 밝혀졌습니다. 2단계 추론 작업에서 첫 번째 루프는 종종 올바른 연결 엔티티를 거의 완벽하게 디코딩하지만, 해당 은닉 상태는 여전히 연결 토큰 임베딩과 제대로 정렬되지 않습니다. 놀랍게도 간단한 학습이 필요 없는 재정렬 방법을 적용하면 일반화 격차를 거의 해소할 수 있습니다. 이러한 통찰력을 바탕으로, 이산 임베딩 채널과 연속적인 은닉 상태 채널을 모두 활용하는 반복적 아키텍처인 DiscoLoop를 제안합니다. DiscoLoop는 기호 및 합성 언어 다중 단계 추론 작업에서 훨씬 적은 학습 단계를 사용하여 거의 완벽한 정확도를 달성합니다. 실제 사전 훈련에 적용했을 때, DiscoLoop는 기존 Looped Transformer 모델보다 낮은 학습 손실과 더 강력한 성능을 보여주며, 이는 혼합 채널 설계가 실용적인 언어 모델링에도 효과적으로 적용될 수 있음을 시사합니다.

Large language models achieve strong performance on many reasoning tasks when allowed to externalize intermediate steps as Chain-of-Thought (CoT). However, many questions require the model to internalize the multi-step reasoning within a single forward pass before generating the answer. We study this challenge through two-hop reasoning, a representative task where the model must compose multiple pieces of parametric knowledge within a single forward pass. Standard non-recurrent Transformers suffer from a depth-local storage problem: facts learned in earlier layers are unavailable where second-hop retrieval happens. We found that Looped Transformers mitigate this issue by reusing the same memory, but still generalize imperfectly. We show that the remaining bottleneck is representational. In the two-hop reasoning task, the first loop often makes the correct bridge entity nearly perfectly decodable, yet the corresponding hidden state remains poorly aligned with the bridge token embedding. Surprisingly, an easy training-free realignment intervention nearly closes the generalization gap. Building upon this insight, we propose DiscoLoop, a looping architecture whose recurrence carries both a discrete embedding channel and a continuous hidden-state channel. DiscoLoop achieves near-perfect accuracy with substantially fewer training steps across symbolic and synthetic-language multi-hop reasoning tasks. When applied to real-world pretraining, DiscoLoop attains lower training loss and stronger benchmark performance than looped-transformer baselines, suggesting that the mixed-channel design transfers to practical language modeling.