[논문 리뷰] Adaptive Loops and Memory in Transformers: Think Harder or Know More?

Markus Frey, Behzad Shomali외 4명|arXiv (Cornell University)|2026. 03. 09.

Topic Modeling인용 수 0

한 줄 요약

이 논문은 적응형 루핑 트랜스포머를 제안하며, 계층별 반복과 게이트된 로컬/글로벌 메모리를 통해 루핑(더 깊이 생각)과 메모리(더 많이 알기)가 수학 및 상식 추론에 어떻게 영향을 주는지 연구합니다. 루핑은 수학 성능을 향상시키고, 메모리는 상식에 도움을 주며, 두 가지를 결합하면 iso-FLOP 기준선보다 더 강력한 성능을 보입니다.

ABSTRACT

Chain-of-thought (CoT) prompting enables reasoning in language models but requires explicit verbalization of intermediate steps. Looped transformers offer an alternative by iteratively refining representations within hidden states. This parameter efficiency comes at a cost, as looped models lack the storage capacity of deeper models which use unique weights per layer. In this work, we investigate transformer models that feature both adaptive per-layer looping, where each transformer block learns to iterate its hidden state via a learned halting mechanism, and gated memory banks, that provide additional learned storage. We find that looping primarily benefits mathematical reasoning, while memory banks help recover performance on commonsense tasks compared to parameter and FLOP matched models. Combining both mechanisms yields a model that outperforms an iso-FLOP baseline, with three times the number of layers, across math benchmarks. Analysis of model internals reveals layer specialization: early layers learn to loop minimally and access memory sparingly, while later layers do both more heavily.

연구 동기 및 목표

• 적응형 계층별 루핑이 트랜스포머의 추론 능력을 향상시키는지 조사합니다.
• 게이트된 메모리 뱅크가 루프형 모델에서 누락된 저장 용량을 회복할 수 있는지 확인합니다.
• 루핑과 메모리가 수학 및 상식 작업에 어떤 상호 작용을 갖는지 분석합니다.
• 루프와 메모리를 사용할 때 내부 계층 특화가 어떻게 나타나는지 특징화합니다.

제안 방법

• 적응형 루핑 및 로컬/글로벌 메모리 뱅크로 디코더-전용 트랜스포머를 확장합니다.
• 학습된 중단 메커니즘을 사용하여 최대 Nmax 반복에 걸친 중간 루프 상태에 가중치를 부여합니다.
• 게이트된 메모리 통합과 함께 로컬(계층별) 및 글로벌(공유) 메모리 KV 뱅크를 도입합니다.
• 약 200M 매개변수 모델의 12층으로 학습하고 14B 토크로 상용 및 수학 벤치마크를 BPB 및 정확도로 평가합니다.
• 루프 증가와 깊이 증가의 효율성을 평가하기 위해 iso-매개변수 및 iso-FLOP 기준선과 비교합니다.

실험 결과

연구 질문

• RQ1 적응형 계층별 루핑이 표준 트랜스포머에 비해 수학적 추론 작업을 개선합니까?
• RQ2 특히 상식 작업에서 루프형 모델의 저장 용량 감소를 메모리 뱅크가 보충합니까?
• RQ3 성과 및 내부 역학 측면에서 루핑 깊이와 메모리 증가 간의 상호 작용은 어떠합니까?
• RQ4 초기 계층과 후기 계층은 루프와 메모리 사용에 어떤 차이를 보합니까?

주요 결과

모델	CS 정확도	CS BPB	수학 BPB
IsoPar	0.477	0.859	2.163
Loop-3	0.501	0.813	1.687
Loop-5	0.503	0.823	1.737
Loop-7	0.498	0.832	1.659
IsoFLOP	0.523	0.780	1.801
IsoPar-M	0.459	0.823	2.108
Mem (g0=-3)	0.472	0.810	1.619
Mem (g0=0)	0.481	0.810	1.662
Mem (g0=3)	0.511	0.794	1.616
IsoFLOP-M	0.535	0.749	1.761

• 적응형 루핑(Nmax=3)은 수학 BPB를 약 0.476 포인트 향상시켜(2.163에서 1.687로) 상식 정확도도 약간 향상시킵니다(0.477에서 0.501으로).
• 루프를 늘리면 수학에서 수익이 감소하는 한계가 있으며, 추가 반복으로 상식 성능이 다소 하락할 수 있습니다.
• 로컬/글로벌 메모리를 추가하면 루핑만으로보다 수학과 상식을 모두 더 개선하며 보완적 이점을 시사합니다.
• 메모리 보강 루핑은 종종 수학 과제에서 iso-매개변수 기준선보다 우수하고 깊이 증가 기반선 대비 상식 격차를 줄입니다.
• 계층 분석은 초기 계층은 루핑을 최소로 하고 메모리에 덜 의존하는 반면, 후기 계층은 더 많이 루프하고 메모리를 더 사용하여 생각하기와 알고 있음을 분화시키는 것을 보여줍니다.
• 루프와 메모리의 조합은 세 배 적은 계층으로도 iso-FLOP 기준선을 능가하며 수학 벤치마크에서 경쟁력 있는 성능을 달성합니다.