LoopCD: 루프 AI의 첫 예측과 마지막 예측을 대조해 더 나은 답 고르기

배경과 문제

루프 트랜스포머는 같은 층 블록을 여러 번 반복 실행하는 구조다. 파라미터 수를 늘리지 않고도 반복 횟수만큼 계산 깊이를 늘릴 수 있어, Universal Transformer 이후 꾸준히 연구되어 왔다. 최근에는 Ouro, Huginn, Parcae처럼 이 구조로 사전 학습한 언어 모델이 공개되었고, 이미 학습된 모델의 중간 층 블록을 학습 없이 반복시키는 Training-Free Looped Transformers도 나왔다. 이 논문이 실험한 Looped-Qwen3는 이 방식을 Qwen3-4B에 적용한 모델이다.

루프 트랜스포머는 반복이 한 번 끝날 때마다 같은 표현 공간 위의 은닉 상태를 만든다. 이 상태를 출력 층에 넣으면 같은 다음 토큰에 대한 예측 분포를 얻을 수 있다. 그러나 일반적인 디코딩은 마지막 반복의 상태만 쓰고 나머지는 버린다. 어려운 문제를 풀 때 처음 떠오른 답과 오래 생각한 끝에 고른 답을 나란히 놓아 보면 생각이 어느 쪽으로 움직였는지가 보이는데, LoopCD는 이 움직임의 방향을 디코딩에 쓰는 방법이다.

기존 접근의 한계

대조 디코딩은 강한 모델의 예측을 약한 모델의 예측에서 멀어지는 방향으로 밀어 주는 방법이다. 두 모델이 함께 높게 평가하는 토큰보다, 강한 모델만 높게 평가하는 토큰을 더 선호하게 만든다. 확산 모델의 분류기 없는 가이던스와 Autoguidance도 같은 원리로, 조건이 없거나 성능을 낮춘 모델의 출력에서 멀어지는 방향으로 샘플을 유도한다.

이 방법은 약한 모델을 어디서 구하느냐가 문제다. 첫 번째 방법은 같은 계열의 작은 모델을 따로 두는 것인데, 모델을 두 개 실행해야 하고 두 모델의 어휘와 행동이 잘 맞아야 한다. 두 번째 방법은 문맥 인식 디코딩처럼 입력 문맥을 뺀 프롬프트로 약한 예측을 만드는 것인데, 같은 입력을 두 가지 형태로 처리해야 한다. 세 번째로 DoLa는 같은 모델의 중간 층을 출력 층에 바로 연결해 약한 예측으로 쓴다. 별도 모델은 필요 없지만, Apple 연구팀은 이렇게 층 방향으로 대조하면 약한 예측과 마지막 예측 사이에 구조적 불일치가 남는다고 본다.

연구방법을 단계별로

연구팀은 루프 트랜스포머라면 별도의 약한 모델 없이, 모델 안에서 약한 참조를 얻을 수 있다고 본다. 첫 반복을 마친 상태와 마지막 반복을 마친 상태는 같은 블록이 만든 같은 표현 공간의 상태이고, 같은 출력 층으로 해석된다. 두 상태는 같은 입력에 대해 계산량만 다르게 들인 결과다. 연구팀의 표현을 빌리면 반복 경로는 '보조 모델이나 외부 학습 없이 정렬된 약한 예측과 강한 예측의 쌍'을 만들어 낸다. 다만 항상 첫 반복을 쓰는 것은 아니다. 초기 잡음의 영향을 줄이기 위해 Huginn의 은닉 상태 변형은 더 뒤의 상태를 참조하며, 반복을 절반으로 줄인 일부 설정도 별도의 참조 상태를 쓴다.

LoopCD는 마지막 반복의 예측을 첫 반복의 예측에서 멀어지는 방향으로 한 번 더 밀어 준다. 반복이 진행되며 예측이 움직인 방향을 그대로 조금 더 연장하는 셈이다. 대조를 어디에서 하느냐에 따라 두 가지 변형이 있다. LoopCD-Logits는 두 상태를 각각 뒤쪽 층과 헤드에 통과시켜 로짓을 얻은 뒤 로짓 공간에서 결합한다. LoopCD-Hidden은 출력 층에 넣기 전에 은닉 상태 공간에서 결합한다.

가이던스 강도는 고정 방식과 적응 방식이 있다. 고정 강도는 한 평가 설정 안의 모든 토큰에 같은 값을 쓴다. 적응 강도는 토큰마다 일반 예측의 불확실성에 맞춰 강도를 조절한다. 상위 두 후보의 확률이 비슷할수록 최대 강도에 가까워지고, 한 후보가 확실히 앞서면 가이던스가 약해진다. 적응 강도를 쓰려면 일반 예측이 먼저 있어야 하므로, 연구팀은 전체 반복 깊이의 LoopCD-Logits에만 이 규칙을 적용했다. LoopCD-Hidden과 반복 절반 실험에는 고정 강도를 썼다.

결과와 숫자의 의미

수학 추론에서 두 Ouro-Thinking 모델 모두 세 벤치마크의 pass@1과 pass@10이 올랐다. Ouro-2.6B-Thinking의 AIME 2024 pass@1은 일반 디코딩 61.88%에서 적응 LoopCD-Logits로 73.33%가 되어 11.45%포인트 올랐다. pass@10의 평균 변화는 +2.53에서 +4.12%포인트였다. 같은 실험을 Looped-Qwen3에 적용한 결과는 부록에 따로 실려 있는데, 여기서는 pass@1이 오히려 떨어졌다. 고정 강도에서 AIME 2024 pass@1은 64.79%에서 61.88%로 내려갔고, 같은 조건의 pass@10은 82.91%에서 86.17%로 올랐다. 수학 평가는 AIME 각 연도의 30문제와 OlympiadBench 영어 수학 674문제를 대상으로, 문제당 16개 풀이를 샘플링해 pass@1과 pass@10을 추정했다. 코드 평가는 이와 달리 그리디 생성과 EvalPlus 테스트를 사용했다.

코드 생성에서는 모델 계열과 크기에 따라 향상 폭이 달랐다. Huginn-0125 R=32에서 LoopCD-Hidden은 HumanEval pass@1을 22.56%에서 31.71%로 올려 9.15%포인트 향상했다. 이는 출력 층을 한 번만 실행하면서도 두 로짓 규칙보다 평균 향상이 컸다. 객관식 7종에서는 6개 모델 모두 평균이 올랐지만, 향상 폭은 1%포인트 안팎이다. 벤치마크별로 보면 ARC-Challenge와 HellaSwag의 향상이 크고, WinoGrande와 PIQA는 여러 설정에서 기준선보다 낮아졌다.

반복 횟수를 절반으로 줄이면 일반 디코딩의 7종 평균은 0.17에서 1.29%포인트 떨어졌다. 같은 절반 깊이에서 LoopCD를 적용하면 6개 설정 모두 전체 반복의 일반 디코딩과 같거나 높아졌다. LoopCD-Logits의 추가 출력 층 비용까지 포함한 순방향 FLOPs는 전체 반복의 0.52에서 0.78배였다. 줄일 수 있는 연산량은 전체 연산에서 반복 블록이 차지하는 비율에 따라 정해진다. 반복 블록의 비중이 큰 Huginn은 46%에서 48%를 줄였다. 뒤쪽 층이 많은 Looped-Qwen3는 LoopCD-Logits를 쓰면 출력 층 추가 비용이 반복 절감분보다 커져서 LoopCD-Hidden으로만 실험했고, 23.6%를 줄였다.

구체적인 예시

논문은 ARC-Challenge 문항 하나로 가이던스가 어떻게 작동하는지 보여 준다. '숲에 좁은 파괴 경로를 남기는 자연재해'를 묻는 문항에서 Ouro-1.4B의 마지막 예측은 'earthquake'(지진, -0.201)를 정답 'tornado'(토네이도, -0.209)보다 0.01 미만 차이로 앞세웠다. 그러나 첫 반복부터 마지막 반복까지 tornado의 점수는 +0.90, earthquake는 +0.32 올랐다. LoopCD는 이 움직임을 연장해 tornado(-0.170)를 earthquake(-0.199)보다 앞으로 올렸다. 연구팀은 가이던스가 새로운 지식을 더하는 것이 아니라 '참조 상태에서 읽어 낸 움직임의 방향을 이어 갈 뿐'이라고 설명한다.

이득이 어디서 나오는지 분석한 결과, 참조 예측이 마지막 예측과 얼마나 다른지가 중요했다. ARC-Challenge에서 Huginn의 첫 반복은 마지막 예측과 다른 답을 고른 비율이 51.3%였고 향상 폭은 +3.07%포인트였다. 16번째 반복은 불일치가 7.3%로 줄어 향상 폭도 +0.17%포인트에 그쳤다. Parcae-1.3B의 6번째 반복은 단독 정확도(41.04%)가 마지막 예측(40.36%)보다 높은데도 향상 폭이 1%포인트에 그쳤다. 이미 마지막 상태에 거의 수렴해서 연장할 방향이 남아 있지 않기 때문이다.

한계와 오해하기 쉬운 점

모든 모델·문제에서 일관 향상이라는 초록 표현을 일반화하면 안 된다. Looped-Qwen3의 수학 추론 pass@1은 떨어졌고, GSM8K는 10개 행 중 6개 행에서만 올랐다. Huginn R=32 고정 강도에서 -1.29%포인트, Looped-Qwen3에서 -1.14%포인트 떨어졌다. 객관식 벤치마크에서도 WinoGrande와 PIQA는 여러 설정에서 기준선보다 낮아졌다.

반복 절반·FLOPs 22.5~48.2% 절감은 객관식 평가 6개 설정에서의 이론적 512토큰 prefill 계산량이다. 실제 초 단위 지연 또는 모든 생성 과제에서 절반 반복 품질 보장으로 해석하면 안 된다. 추가 학습이나 외부 보조 모델은 없지만 logits 변형은 출력층 추가 실행 비용이 들며 Looped-Qwen3에서 연산 1.306배다. Hidden의 추가 출력층 연산은 없다.

검증 범위가 작은 모델에 한정된다. 실험한 모델은 0.37B에서 4B 사이다. 대형 모델에서 같은 효과가 나는지는 확인되지 않았다. 강도 선택에 별도 검증 세트가 언급되지 않는다. 연구팀은 가이던스 강도와 참조 반복을 본 평가 전 스크리닝 스윕으로 골랐다고 밝혔는데, 논문이 보여 주는 스윕 곡선은 평가에 쓴 같은 벤치마크의 평균이다. 객관식 향상 폭이 1%포인트 안팎이라는 점을 함께 고려하면, 새 과제에 적용할 때는 강도를 다시 찾아야 할 가능성이 있다.

실제로 어떻게 활용할까

루프 트랜스포머를 서빙하는 환경에서 LoopCD-Hidden은 출력 층 연산을 늘리지 않으면서 정확도를 높일 수 있는 후보 방법이다. 특히 Huginn처럼 반복 블록의 비중이 큰 모델에서는 객관식 7종 평균에서 반복을 절반으로 줄여도 전체 반복의 일반 디코딩과 같거나 높은 정확도를 보였다. 이 결과를 서비스의 생성 품질이나 실제 지연 시간에 적용하려면 별도의 측정이 필요하다.

새 과제에 적용할 때는 가이던스 강도와 참조 반복을 별도의 검증 자료로 선택하는 것이 합리적이다. 논문의 주요 생성 설정은 강도 0.2~0.3, 객관식은 0.5 근처를 사용했지만, 이는 평가한 설정의 선택값이며 모든 모델에 의무적으로 적용할 값은 아니다. Huginn의 은닉 상태 실험은 초기 잡음을 피하려고 여섯째 또는 일곱째 상태를 참조했다. 한국어 해설이 소개하는 공개 저장소는 이름이 같은 별도 POSTECH 연구의 공식 구현이다. 이를 Apple 연구 전체의 구현으로 혼동해서는 안 되며, 해설에 따르면 해당 저장소는 Huginn-0125와 Ouro-1.4B를 지원하고 Apple의 은닉 상태 변형과 적응 강도는 포함하지 않는다.