💬 NLP

LoopMTP: A looped transformer guided by latent multi-token prediction

이 논문은 중간 은닉 상태를 잠재적 다중 토큰 예측을 통해 미래 토큰 임베딩과 정렬함으로써 과잉 사고(overthinking)를 완화하고 추론 정확도를 크게 향상시키는 방식으로 루프드 트랜스포머(looped transformers)를 개선하는 매개변수 효율적인 프레임워크인 LoopMTP를 제안한다.

원저자: Behzad Shomali, Markus Frey, David Berghaus, Joachim Koehler, Mehdi Ali

게시일 2026-08-05

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Behzad Shomali, Markus Frey, David Berghaus, Joachim Koehler, Mehdi Ali

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 아주 까다로운 퍼즐, 예를 들어 복잡한 수학 문제나 반전이 있는 이야기를 쓰는 일을 하고 있다고 상상해 보세요. 당신은 각자의 책상을 가진 전문가들로 구성된 거대한 팀을 고용할 수도 있을 것입니다. 이것이 오늘날 대부분의 강력한 AI 모델이 작동하는 방식입니다. 즉, 수천 개의 '전문가' 층이 서로 쌓여 있는 거대하고 거대한 구조를 가지고 있습니다. 하지만 만약 당신에게 단 한 명의 명석한 전문가와 작은 책상 하나뿐이라면 어떨까요? 그 전문가에게 문제를 계속해서 반복해서 생각하며 답변을 정교하게 다듬을 기회를 준다면, 여전히 그 퍼즐을 풀 수 있을까요? 이것이 바로 '루프 트랜스포머(looped transformers)'라고 불리는 분야에서 논의되고 있는 핵심적인 질문입니다.

오랫동안 과학자들은 AI가 더 똑똑해지기 위해서는 더 커져야 한다고 믿어 왔습니다. 하지만 여기에는 함정이 있습니다. 모델이 커질수록 비용이 많이 들고 일반 컴퓨터에서 실행하기 어렵다는 점입니다. 또한, AI는 때때로 자신의 생각 속에 갇혀 단순한 답을 너무 깊게 생각하다가 오히려 망쳐버리는 루프에 빠지기도 합니다. 또 다른 아이디어인 '멀티 토큰 예측(Multi-Token Prediction)'은 AI가 단순히 문장의 다음 단어를 추측하는 대신, 동시에 다음 몇 개의 단어를 예측해야 한다고 제안합니다. 이는 마치 체스 선수가 단 한 수 앞이 아니라 세 수 앞을 내다보는 것처럼, AI가 앞을 내다보며 계획을 세우도록 강제합니다. 과제는 '다시 생각하는' 접근 방식과 '앞을 내end다보는' 접근 방식을 어떻게 결합하여 AI가 혼란에 빠지거나 에너지를 낭비하지 않게 하느냐였습니다.

여기에 연구자 베자드 쇼말리(Behzad Shomali)와 그의 팀이 제안한 새로운 방법인 LOOPMTP가 등장했습니다. LOOPMTP를 이 단 한 명의 명석한 전문가를 위한 영리한 전략이라고 생각해 보세요. 전문가가 퍼즐을 뚫어지게 쳐다보며 노트를 다시 쓰는 대신(이는 종종 좋은 아이디어를 지워버리는 결과를 초래합니다), LOOPMMTP는 그들에게 특별한 가이드를 제공합니다. 전문가가 퍼즐을 새로 살펴보는 매 순간(즉, '루프'를 돌 때마다), 그들이 해결책의 특정 미래 조각에 시선을 고정하도록 부드럽게 유도하는 것입니다.

이것이 실제로 어떻게 작동하는지 살펴보겠습니다. AI가 이야기를 쓰고 있다고 가정해 봅시다. 표준적인 루프에서는 AI가 문장을 쓴 다음, 방금 쓴 내용을 즉시 잊어버리고 그것을 다시 쓰려고 시도하다가 종종 실수를 저지릅니다. LOOPMTP는 규칙을 바꿉니다. AI가 이야기의 두 번째 패스를 진행할 때, 내부적으로 "헤이, 지금부터 두 단계 뒤에 나올 단어를 준비해야 한다는 걸 기억해"라는 메시지를 받게 됩니다. 세 번째 패스에서는 세 단계 앞을 내다보라는 지침을 받습니다. 이것이 '멀티 토큰 예측'이 가이드 역할을 하는 부분입니다. 이는 AI에게 미래의 단어들을 소리 높여 외치라고 강요하는 것이 아니라, 그 미래 단어들의 '개념'에 AI의 내부적인 생각을 정렬시키는 것입니다. 이는 마치 등산객이 앞으로 달려 나가기 위해서가 아니라, 매 걸음이 올바른 방향으로 향하고 있는지 확인하기 위해 주머니에 지도를 넣어두는 것과 같습니다.

연구진은 또한 '문지기(gatekeeper)' 메커니즘을 추가했습니다. 과거에는 AI가 루프를 돌 때, 마치 요리사가 맛을 볼 때마다 소스를 버리는 것처럼 이전 루프의 작업물을 버리곤 했습니다. LOOPMTP는 모든 버전의 소스를 보관하면서, 스마트한 게이트를 사용하여 각 버전을 최종 요리에 얼마나 섞을지 결정합니다. 이를 통해 AI가 문제를 여러 번 생각하더라도 좋은 정보가 손실되지 않도록 보장합니다.

이러적인 접근 방식의 결과는 매우 유망합니다. 연구팀이 수학, 코딩, 일반 지식 질문을 포함한 다양한 과업에 대해 LOOPMTP를 테스트했을 때, 이 모델은 루핑 기술을 사용하지 않은 표준 모델보다 성능이 현저히 뛰어났습니다. 실제로, 이들은 루핑 기술을 사용하지 않은 표준 모델과 비교했을 때 평균적으로 최대 **8.1%**까지 정확도를 향상시켰습니다. 더욱 인상적인 점은, AI에게 문제를 15번 연속으로 "생각"하도록 요청했을 때도 안정성을 유지하며 계속해서 개선되는 모습을 보였다는 것입니다. 유사한 다른 모델들은 불과 몇 번의 루프만으로도 무너지거나 혼란에 빠졌습니다.

이 논문은 이 방법이 거대하고 비싼 모델을 구축할 필요 없이 더 작은 AI 모델을 더 똑똑하게 만드는 강력한 방법임을 시사합니다. AI의 내부적인 생각을 미래로 유도함으로써, 우리는 AI가 과하게 생각하는 것을 방지하고 제한된 컴퓨팅 자원을 훨씬 더 효율적으로 사용하도록 도울 수 있다는 것을 보여줍니다. 연구진은 이 방법이 특히 수학 및 추론 과업에 효과적이라고 언급하면서도, 한계가 있다는 점 또한 지적했습니다. 즉, 모델에게 너무 많이 생각하도록 요구하면 결국 다시 혼란을 겪을 수 있다는 것입니다. 하지만 현재로서 LOOPMTP는 AI가 더 큰 뇌를 갖지 않고도 더 열심히 "생각"할 수 있도록 돕는 신선하고 즐거운 방식을 제시하고 있습니다.

기술 요약: LoopMTP

문제 정의
대규모 언어 모델(LLM)은 강력한 추론 능력을 입증해 왔으며, 이는 종종 아키텍처의 깊이와 연관되어 있습니다. 그러나 깊이가 증가할수록 상당한 파라미터 수가 필요하며, 이는 자원이 제한된 환경(예: 온프레미스 의료 배포)에서는 비현실적입니다. 루프형 트랜스포머(Looped transformers)는 단일 레이어 스택을 여러 번 반복( $T$ )하여 재사용함으로써 유효 깊이를 시뮬레이션하는 파라미터 효율적인 대안을 제공합니다. 이러한 잠재력에도 불구하고, 기존의 루프 모델들은 두 가지 결정적인 한계점을 가집니다:

잠재적 과잉 사고(Latent Overthinking): 중간 표현들이 각 반복 과정에서 덮어씌워집니다. 만약 예측이 초기에 이미 정확했다면, 후속 루프가 이를 오류로 수정해 버릴 수 있습니다.
차별화되지 않은 계산(Undifferentiated Computation): 반복 횟수가 증가함에 따라 은닉 상태(hidden states)가 점점 더 유사해지며, 이는 중복된 작업과 계산 낭비를 초래합니다.
또한, 기존 방식들은 루프 전반에 걸친 가이드(guidance)가 부족하며, 미래를 내다보는 추론을 장려하는 다중 토큰 예측(Multi-Token Prediction, MTP)은 역사적으로 대규모 규모나 비용이 많이 드는 전체 어휘 투영(full-vocabulary projections)이 동반될 때만 이점이 있는 것으로 나타났습니다.

방법론: LoopMTP
저자들은 과잉 사고와 차별화되지 않은 계산 문제를 해결하기 위해, MTP를 잠재적 가이드 신호로 통합한 루프형 트랜스포머 아키텍처인 LoopMTP를 제안합니다. 핵심 구성 요소는 다음과 같습니다:

MTP 가이드 루프 블록 (MTP-Guided Looped Block): 모델은 전체 트랜스포머 스택을 $T$ 번 루프합니다. 유용한 정보가 덮어씌워지는 것을 방지하고 궤적을 가이드하기 위해, 아키텍처는 다음을 채택합니다:
- 반복 인덱스 임베딩 (Iteration-Index Embedding): 정규화된 반복 인덱스를 토큰 벡터에 추가하여 현재 단계를 알립니다.
- 반복별 정규화 (Per-Iteration Normalization): 반복별 통계치를 처리하기 위해 각 반복의 상태에 특화된 LayerNorm을 적용합니다.
- 입력 퓨전 (Input Fusion): 이전 반복의 출력과 재정규화된 토크 임베딩을 결합하고 투영(project)한 후 공유 백본(backbone)으로 진입합니다.
- Loop-LNS (Looped Layer-Norm Scaling): 가변적인 언롤링(unrolling) 상황에서도 안정성을 유지하고, 점진적 깊이 확장 시 관찰되는 그래디언트 중립화를 피하기 위해 잔차 연결(residual streams)에 $1/T$ 의 수정된 스케일링 계수를 적용합니다.
소프트 MTP 정렬 (Soft MTP Alignment, 잠재 가이드): 전체 토큰 분포를 예측하는 대신(이는 계산 비용이 많이 듭니다), LoopMTP는 소프트 정렬 목적 함수를 사용합니다. 반복 $t$ (단, $t \ge 2$ )에 대해, 은닉 상태는 $t$ 스텝 앞의 토큰 임베딩( $u_{i+t}$ )과의 코사인 유사도를 통해 정렬됩니다. 이는 표현의 정규화 도구 역할을 하며, 전체 어휘 투영 없이도 각 반복이 특정 미래 토큰을 예측하도록 유도합니다. 첫 번째 반복( $t=1$ )은 풍부한 초기 표현을 학습할 수 있도록 제약을 두지 않습니다.
학습된 상태 집계 (Learned State Aggregation): 중간 상태를 버리는 대신, LoopMTP는 모든 $T$ 번의 반복 출력을 토큰 단위 가중 합(weighted sum)을 사용하여 집계합니다. 경량화된 공유 선형 게이트(linear gate)가 반복별 가중치를 계산하여, 모델이 초기 루프의 유용한 정보를 유지하면서 후기 루프의 새로운 계산을 통합할 수 있게 합니다. "폰더 정규화 도구(ponder regularizer, KL 발산)"는 반복 간의 주의(attention)가 균등하게 분포되도록 유도하여 조기 수렴을 방지합니다.
학습 목적 함수 (Training Objective): 총 손실(total loss)은 표준 차기 토큰 예측(NTP) 교차 엔트로피, 소프트 MTP 정렬 손실, 그리고 폰더 정규화 도구를 결합합니다.

주요 기여

아키텍처 혁신: 루프의 수를 미래 토큰의 예측과 구조적으로 연결하여, 잠재적 과잉 사고와 차별화되지 않은 계산에 대한 자연스러운 해결책을 제공하는 LoopMTP를 도입했습니다.
효율적인 MTP 통합: 경량화된 코사인 정렬 신호를 통해 MTP가 파라미터 효율적인 루프형 모델에 효과적으로 적용될 수 있음을 입증하였으며, 이를 통해 표준 MTP의 계산 오버헤드를 피했습니다.
안정적인 학습: 제안된 아키텍처는 LoopFormer와 같은 기존 루프 모델에서 흔히 발생하는 문제인, 루프 횟수당 광범위한 하이퍼파라미터 그리드 탐색 없이도 넓은 범위의 루프 횟수(최대 $T=15$ )에 대해 안정적인 학습을 가능하게 합니다.
소형 전문가 모델: LoopMTP가 동일한 파라미터 예산 하에서 더 큰 비루프형(non-looped) 베이스라인보다 뛰어난 성능을 보이는 소규모 도메인 특화 전문가 모델(예: 수학용)을 학습시킬 수 있음을 보여주었습니다.

결과
6.8B 토큰으로 학습된 260M 파라미터 모델을 대상으로 평가한 결과:

성능 향상: LoopMTP는 일반 작업에서 파라미터가 일치하는 비루프형 베이스라인 대비 상대적 평균 정확도에서 최대 8.1% 상회하는 성능을 보였습니다. 또한 QA, 수학, 코드 작업에서 비루프형 베이스라인 대비 최대 7.5% 낮은 Bits-Per-Byte(BPB)를 기록했으며, SOTA인 LoopFormer를 크게 앞질렀습니다 (28개의 매칭된 비교 중 27개에서 승리).
확장성: 높은 루프 횟수에서 불안정한 학습과 성능 저하를 겪는 LoopFormer와 달리, LoopMTP는 $T$ 가 증가함에 따라 퍼플렉서티(perplexity)와 일반 작업 정확도에서 단조로운 개선을 보여줍니다.
메커니즘 분석:
- 표현(Representations): MTP를 적용했을 때 반복 간의 은닉 상태가 더욱 뚜렷해졌으며, 이는 모델이 중복된 것이 아닌 고유한 변환을 학습하고 있음을 확인시켜 줍니다.
- 정답 추출(Ground-Truth Retrieval): MTP 가이드가 있을 때 정답 토큰 순위(rank)가 최대 35.6배 유의미하게 개선되었으며, 이는 정렬 신호가 언어 모델 헤드로 효과적으로 전달됨을 나타냅니다.
- 게이팅(Gating): 학습된 게이팅 메커니즘은 깊이에 따라 적응하며, 후기 반복에서 더 높고 가변적인 게이트 값을 보여 정보 유지를 위한 입력 의존적 전략을 나타냅니다.
도메인 전문가: 특화된 수학 학습 설정에서, $T=15$ 인 LoopMTP 모델은 GSM8K에서 **19.03%**의 정확도를 달 기록하였으며, 이는 동일 파라미터 수의 비루프형 베이스라인(7.05%) 대비 약 170%의 상대적 개선입니다.

의의 및 주장
본 논문은 LoopMTP가 파라미터 효율성과 강력한 추론 능력 사이의 간극을 성공적으로 메웠다고 주장합니다. MTP를 잠재적 가이드 신호로 활용함으로써, 이전의 루프형 트랜스포머가 가진 근본적인 불안정성과 중복 문제를 해결했습니다. 저자들은 이 접근 방식이 데이터가 외부로 유출될 수 없는 의료와 같은 민감한 분야에서 필수적인, 제한된 하드웨어에서의 강력한 추론 모델 배포를 가능하게 한다고 강조합니다. 이 연구는 미래를 내다보는 목적 함수에 의해 가이드된 "유효 깊이(effective depth)"가 루프를 통한 데 깊이가 원시적인 파라미터 스케일링보다 특정 추론 작업에서 더 나은 성능을 낼 수 있음을 시사합니다. 저자들은 수학 이외의 특화된 도메인에서의 이점은 추가 검증이 필요하며, 유용한 루프 깊이의 상한선은 여전히 열린 문제라는 점을 언급하며 신중한 태도를 유지하고 있습니다.

기술 요약: LoopMTP

유사한 논문