ARLoop는 반복 깊이에 어텐션 잔차를 적용해 루프 트랜스포머를 강화한다
ARLoop: Fully-Open Pre-trained Looped Transformers with Attention Residuals
무엇인가
루프 트랜스포머는 같은 트랜스포머 층 묶음을 여러 번 재사용해 파라미터 수를 늘리지 않고 유효 깊이를 키우는 구조다. 문제는 표준 PreNorm 잔차 연결이 이전 깊이에서 나온 모든 잔차 업데이트를 고정된 단위 계수로 하나의 스트림에 누적한다는 점이다. 반복 깊이가 늘어날수록 새 중간 표현이 계속 생기지만, 뒤쪽 계산은 이전 깊이의 어떤 표현을 재사용할지 명시적으로 고를 수 없다. 논문은 이 지점을 루프 트랜스포머가 추가 계산을 제대로 활용하지 못하는 핵심 병목으로 본다.
어떻게 동작하나
ARLoop는 Kimi Team이 제안한 어텐션 잔차를 반복 깊이에 맞게 변형한다. 각 하위층은 z_j = f_j(Norm(h_{j-1}))로 잔차 업데이트를 만들고, 학습된 질의 q_i로 이전 잔차 소스들에 어텐션을 걸어 alpha_i,j를 계산한 뒤 h_i = alpha_i,0 h_0 + sum alpha_i,j z_j로 입력을 선택적으로 집계한다. 즉 PreNorm처럼 모든 업데이트를 균일하게 더하지 않고, 깊이 방향 어텐션으로 필요한 표현의 비중을 조절한다. 잔차 소스가 깊이에 따라 선형으로 늘어나는 문제는 연속된 하위층을 블록으로 묶는 블록 어텐션 잔차로 줄인다. 블록 표현은 b_k = sum z_j로 만들고, 각 하위층은 입력 임베딩, 이전 완료 블록, 현재 블록의 부분 합에 어텐션을 수행한다.
무엇과 다른가
논문은 루프 트랜스포머에서는 전체 잔차 이력이 필요하지 않을 수 있다고 주장한다. 같은 파라미터를 반복 적용하면 반복 깊이 사이 표현이 중복될 수 있기 때문이다. 그래서 입력 임베딩과 가장 최근 W개 잔차 블록만 보는 슬라이딩 윈도 깊이 어텐션을 비교한다. 실험에서는 R이 2 이하일 때 전체 이력 어텐션을, R이 2보다 클 때 슬라이딩 윈도를 기본으로 쓰고 최근 블록 3개를 유지한다. 분석에서 깊이 어텐션은 최근 잔차 소스에 집중됐고, 47번째 하위층에서는 최근 블록 3개와 임베딩 블록이 전체 어텐션 가중치의 96%를 받았다. 반복 깊이 사이 잔차 블록의 코사인 유사도도 높아 중복성이 크다는 설명이 붙는다.
어떻게 쓰나
통제 실험은 140M과 370M 파라미터의 decoder-only 모델로 진행됐다. 각각 6층과 12층을 쓰고 FineWeb-Edu에서 11.2B와 29.6B 토큰으로 사전학습했으며 컨텍스트 길이는 2048이다. 비교 대상은 PreNormLoop, Sandwich Norm, DeepLoop, Parcae와 입력 주입 변형이다. 140M에서는 R=1,3,5를, 370M에서는 R=1,3을 비교했다. 평가는 10M 토큰 검증 세트의 퍼플렉시티와 ARC, HellaSwag, LAMBADA, OpenBookQA, PIQA, RACE, SciQ, WinoGrande의 평균이다. ARLoop는 모든 조건에서 가장 낮은 검증 퍼플렉시티를 기록했고, PreNormLoop 대비 개선폭은 R=1의 -0.11 PPL보다 R=3의 -0.53 PPL에서 더 컸다. 슬라이딩 윈도 변형은 140M의 R=2에서는 전체 이력보다 나빴지만, 더 큰 반복 깊이와 큰 파라미터 규모에서는 가장 좋은 성능을 보였다.
전제와 한계
실제 LLM 사전학습 레시피에서는 SmolLM과 유사하게 grouped-query attention과 SwiGLU를 쓰고, FineWeb-Edu, FineMath, Stack-Edu를 주 데이터로 삼았다. 후반 단계에 고품질 데이터를 넣는 다단계 학습을 적용하고 컨텍스트 길이를 4096으로 늘렸다. PreNormLoop-0.6B는 18층, R=3, 180B 토큰으로 학습했고, 같은 조건의 ARLoop-0.6B와 24층, R=3, 500B 토큰으로 학습한 ARLoop-1.2B를 만들었다. ARLoop-0.6B는 8개 과제 중 7개에서 PreNormLoop-0.6B를 이겼고, GSM8K에서 +5.2점, MBPP에서 +3.1점을 냈다. ARLoop-1.2B는 Huginn-3.5B를 모든 평가 과제에서 앞섰으며, 파라미터·토큰·반복 횟수가 더 적음에도 그런 결과를 냈다. Ouro에는 대부분 과제에서 뒤졌지만, 논문은 그 차이를 Ouro의 더 큰 사전학습 데이터 예산과 고품질 데이터 사용으로 설명한다. ARLoop-1.2B는 Qwen3.5-2B와도 모든 벤치마크에서 맞먹거나 앞서고, Gemma-3-4B보다 수학 벤치마크에서, OLMo 모델보다 HumanEval에서 앞선다고 보고한다.
저자들은 ARLoop의 이점이 단순 퍼플렉시티를 넘어 깊은 반복 계산을 요구하는 과제에서 커진다고 본다. layer effect 분석에서 ARLoop-0.6B의 두 번째와 세 번째 반복에 있는 층들이 PreNormLoop-0.6B의 해당 깊은 층들보다 이후 층에 더 큰 영향을 줬다. 또한 동적 반복 학습보다 고정 R=4 학습이 140M과 370M 모두에서 더 나은 검증 퍼플렉시티와 하위 과제 성능을 냈다. 반면 추론 시 학습 때보다 많은 반복을 주는 외삽은 최대 학습 반복까지만 성능이 오르고 그 이상에서는 개선되지 않았다. 논문은 선택적 잔차 접근만으로는 반복 계산 외삽이 열리지 않는다고 명시한다.
개발자 관점에서 이 논문은 루프 트랜스포머의 잔차 연결을 단순한 구현 세부사항으로 보면 안 된다는 신호를 준다. PreNorm 기반 루프 모델은 추가 반복 계산을 충분히 활용하지 못할 수 있고, 블록 단위 어텐션 잔차와 슬라이딩 윈도는 비교적 적은 메모리로 이를 개선할 수 있는 선택지다. 다만 슬라이딩 윈도 어텐션 잔차는 PreNorm 대비 처리량이 약 10% 느리고, 전체 이력 어텐션은 약 15% 느리다. 활성화 체크포인팅을 쓰지 않으면 메모리 차이는 거의 없지만, 깊이 어텐션의 커널 최적화 여지가 남아 있다고 저자들은 밝힌다.
한계도 분명하다. 실제 사전학습 비교에서는 Parcae나 DeepLoop 같은 다른 잔차 설계를 포함하지 않았는데, 논문은 그런 연구들이 보통 50B 토큰 이하의 작은 데이터 규모나 단일 데이터 소스에서만 검증돼 수학·코딩 같은 복잡한 생성 과제에서의 효과가 열려 있다고 설명한다. Ouro와의 격차도 데이터 예산과 품질 차이로 해석되며, ARLoop 자체의 우위만으로 단정하기는 어렵다. 또한 추론 시 반복 외삽은 학습된 반복 범위를 넘어서면 개선되지 않았고, 처리량 오버헤드와 커널 최적화 필요성도 남는다.