🤖 AI

Looped State-Space Language Models with Adaptive Exit-State Selection

이 논문은 맘바(Mamba) 기반 상태 공간 모델에 루프형 아키텍처를 적용하는 것이 계산 깊이의 증가를 통해 추론 능력과 파라미터 효율성을 향상시키는 동시에, 예측 깊이를 최적화하기 위한 적응형 종료 상태 선택을 도입하지만, 실제 추론 시의 절감을 위해서는 추가적인 상태 처리 메커니즘이 필요하다고 언급한다.

원저자: Zhenxuan Yu, Takeshi Kojima, Yutaka Matsuo, Yusuke Iwasawa

게시일 2026-07-14

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Zhenxuan Yu, Takeshi Kojima, Yutaka Matsuo, Yusuke Iwasawa

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

당신에게 아주 똑똑하고 천재적인 로봇 셰프 '맘바(Mamba)'가 있다고 상상해 보세요. 이 셰프는 엄청나게 빠르고 메모리 효율이 높은 것으로 유명하며, 다양한 도구가 가득한 거대한 주방 없이도 폭풍처럼 요리를 해낼 수 있습니다. 하지만 한 가지 문제가 있습니다. 정말 까다로운 레시피(예를 들어 복잡한 수학 퍼즐이나 '숨겨진 재료 찾기' 게임 같은 것)를 마주하면, 맘바는 가끔 막히곤 합니다. 더 깊이 생각해야 하지만, 거대한 주방을 새로 짓지 않고서는(그것은 엄청난 비용이 듭니다) 더 깊은 사고를 할 수 있는 충분한 '레이어(층)'를 가지고 있지 않기 때문입니다.

여기서 릿쿄 대학교와 도쿄 대학교의 연구진이 등장합니다. 그들은 단순하면서도 대담한 질문을 던졌습니다. "만약 우리가 더 큰 주방을 만드는 대신, 맘다가 같은 요리를 반복해서 만들며 매 패스(pass)마다 맛을 정교하게 다듬게 한다면 어떨까?"

"루프형(Looped)" 주방: 한 명의 셰프, 여러 번의 패스

AI의 세계에서 대부분의 모델은 공장의 조립 라인과 같습니다. 데이터(문장 속의 단어 같은 것)가 라인을 따라 이동하며 24개의 서로 다른 스테이션(레이어)을 통과한 뒤 작업이 끝납니다. 만약 로봇이 더 깊이 생각하게 하고 싶다면, 보통은 스테이션을 더 추가합니다.

저자들은 다른 시도를 했습니다. 그들은 **'루프형 맘다(Looped Mamba)'**를 만들었습니다. 맘다가 하나의 초능력을 가진 스테이션이라고 상상해 보세요. 데이터를 24개의 서로 다른 스테이션으로 보내는 대신, 그들은 데이터를 동일한 스테이션으로 24번 연속해서 다시 보냅니다.

마법 같은 점: 로봇은 매 패스마다 정확히 동일한 뇌(파라미터)를 사용합니다. 이는 마치 학생이 네 개의 서로 다른 혼란스러운 문단을 읽는 대신, 이해를 돕기 위해 같은 문단을 네 번 다시 읽는 것과 같습니다.
결과: 이 "다시 읽기" 전략은 Mano(모듈로 연산)와 p-hop induction(숨겨진 패턴 찾기)이라는 까다로운 논리 퍼즐에서 놀라운 효과를 발휘했습니다. 단 4개의 레이어만 있지만 이를 6번 루프시킨 모델(4 ⊗ 6)이 24개의 고유한 레이어를 가진 표준 모델을 압도했습니다. 이는 추론에 있어서 생각의 깊이가 도구의 개수보다 더 중요하다는 것을 시사합니다.

"출구 게이트(Exit Gate)": 최선의 버전 선택하기

하지만 잠깐, 로봇이 영원히 문단을 다시 읽게 된다면 어떻게 될까요? 그것은 시간 낭비입니다. 연구진은 **'출구 게이트(Exit Gate)'**라는 영리한 기능을 추가했습니다.

이 게이트를 스마트한 선택기라고 생각하세요. 맘다가 데이터를 처리함에 따라, 게이트는 질문을 던집니다. "답이 명확해지고 있는가?"

만약 2번의 루프 후에 이미 답이 명확하다면, 게이트는 "멈춰! 이 패스의 결과를 사용하자!"라고 말합니다.
만약 문제가 어렵다면, 게이트는 "계속해!"라고 말하며 3번째 또는 4번째 루프의 결과를 선택합니다.

이것을 **'적응형 종료 상태 선택(Adaptive Exit-State Selection)'**이라고 부릅니다. 논문은 1억 4천만 파라미터 규모의 작은 모델에서 이 게이트가 게임 체인저였다고 밝혔습니다. 이 게이트는 모델이 각 단어에 대해 '딱 적절한' 양의 생각 시간을 선택할 수 있게 해주었으며, 종종 최대 루프 횟수를 강제로 사용하는 것보다 더 나은 성능을 보였습니다.

하지만 반전이 있습니다. 연구진은 이 방식이 아직 컴퓨터의 전기를 아끼거나 시간을 단로 줄여주지는 않는다는 점을 매우 주의 깊게 언급했습니다. 게이트가 더 이른 결과물을 선택하더라도, 로봇은 배경에서 여전히 모든 루프를 물리적으로 실행합니다. 이는 맘다의 메모리 상태가 연속적이기 때문입니다: 루프 3의 결과는 루프 2가 남긴 상태에 의존하고, 루프 2는 루프 1에 의존합니다. 따라서 이후의 모든 단어를 위한 체인을 깨뜨리지 않고 중간에 로봇을 그냥 "내릴" 수는 없습니다. 즉, 게이트는 예측 깊이(어떤 버전의 답을 사용할지)를 선택하지만, *실제 계산 시간(wall-clock computation)*은 동일하게 유지됩니다. 실제로 시간을 절약하려면 맘다의 메모리 상태를 다루는 새로운 방법이 필요한데, 저자들은 이것이 향후 과제라고 말합니다.

현실적인 점검: 작동하지 않은 것 (그리고 여전히 미지수인 것)

이 논문은 자신들이 해결하지 못한 부분에 대해서도 솔직하게 밝히고 있습니다.

이 기술에서 '크다고 항상 좋은 것은 아니다': 연구진은 루프형 맘다를 동일한 비용(동일한 계산량, 즉 "iso-FLOPs")을 사용하는 표준 비루프형 모델과 비교했습니다. 이때 표준 모델은 "모델이 얼마나 혼란스러워하는가?"(Perplexity)라는 기본 테스트에서 여전히 승리했습니다. 루프형 모델은 추론 작업에는 뛰어났지만, 단순히 다음 단어를 예측하는 데 있어서는 표준 딥 모델이 여전히 약간 더 나았습니다.
규모가 중요합니다: "출구 게이트"는 1억 4천만 파라미터 모델에서는 아름답게 작동했습니다. 하지만 3억 7천만 파라미터의 더 큰 모델에 적용했을 때는 게이트가 큰 개선을 보여주지 못했습니다. 더 큰 모델들은 그냥 모든 루프를 실행하는 것을 선호하는 듯 보였습니다. 저자들은 더 큰 모델들이 마지막의 깊은 사고 과정에 더 많이 의존할 수도 있다고 제안합니다.
모든 것에 만능인 마법의 탄환은 아닙니다: 연구진은 이 기술을 수학 트리나 패턴 찾기와 같은 합성 퍼즐(synthetic puzzles)에서 테스트했습니다. 이러한 테스트는 개념이 작동함을 증证明하지만, 이 "루핑" 기술이 소설 쓰기나 복잡한 법률 사례 해결과 같은 개방형의 실제 세계 추론에서도 동일하게 잘 작동할지는 아직 알 수 없다고 인정했습니다.

결론

이 논문은 똑똑한 뇌를 반복해서 재사용하는 것이 더 크고 비싼 뇌를 만들지 않고도 AI를 추론에 더 똑똑하게 만드는 강력한 방법임을 시사합니다. 이는 "파라미터 효율적인(parameter-efficient)" 확장 축입니다.

증명됨: 특정 논리 퍼즐에서, 작은 맘다 모델을 루프시키는 것은 동일한 크기의 표준 모델을 이기며 훨씬 더 큰 모델과 대등한 성능을 보여줍니다.
제안됨: 이 접근 방식은 특히 단계별의 깊은 사고를 요구하는 작업에서 미래의 AI를 더 효율적으로 만드는 핵심 열쇠가 될 수 있습니다.
아직 해결되지 않음: 우리는 "출구 게이트"가 (단순히 출력을 선택하는 것을 넘어) 어떻게 실제로 계산 시간을 절약할 수 있을지, 그리고 이 기술이 오늘날 가장 큰 AI를 움직이는 수십억 개의 파라미터 모델에도 적용될 수 있을지를 여전히 풀어야 합니다.

요약하자면, 저자들은 단순히 더 큰 로봇을 만든 것이 아니라, 작은 로봇이 문제에 대해 두 번(또는 세 번, 혹은 네 번) 더 들여다볼 수 있게 함으로써 더 깊이 생각하도록 가르친 것입니다. 그리고 현재로서는, 그 "두 번째 눈길"은 꽤 유망해 보입니다.

기술 요약: 적응형 종료 상태 선택을 갖춘 루프형 상태 공간 언어 모델 (Looped State-Space Language Models with Adaptive Exit-State Selection)

문제 정의

최근 대규모 언어 모델(LLM)의 발전은 주로 모델 크기, 학습 데이터, 그리고 연산 집약적인 추론의 스케일링에 의해 주도되어 왔습니다. 그러나 이러한 스케일링 패러다임은 Transformer 기반 셀프 어텐션 메커니즘의 거대한 키-값(KV) 캐시 유지로 인한 메모리 소비 증가 및 급격히 늘어나는 추론 비용을 포함한 심각한 배포 과제를 야기합니다. Mamba-2와 같은 상태 공간 모델(SSM)은 유리한 롱 컨텍스트 스케일링을 갖춘 효율적인 대안을 제공하지만, 최근 Looped Transformer에서 입증된 "독립적 파라미터보다 계산적 깊이(computational depth over independent parameters)"의 원리가 SSM에도 적용되는지는 여전히 불분명합니다. 또한, 기존의 루프형 아키텍처 연구는 거의 독점적으로 Transformer 백본에 집중되어 있어, Mamba 및 하이브리드 Mamba-Transformer 아키텍처 내에서 재귀적 파라미터 공유가 어떻게 기능하는지에 대한 이해의 공백이 존재합니다.

방법론

저자들은 Looped Mamba 및 Looped Hybrid Mamba–Transformer 아키텍처를 제안합니다. 핵심 메커니즘은 단일 Mamba(또는 하이브리드) 블록의 파라미터를 여러 레이어에 걸쳐 공유하고, 이를 명시적인 유한 깊이 루프를 통해 반복적으로 적용하는 것입니다.

아키텍처

표기법: 모델은 $N \otimes R$ 로 표기되며, 여기서 $N$ 은 독립적으로 파라미터화된 레이어의 수이고 $R$ 은 재귀 단계(recurrent steps)의 수입니다. 유효 깊이는 $NR$입니다.
Looped Mamba: 공유된 $N$ $N$ -레이어 Mamba 백본 $M_\theta$ $M_{θ}$ 를 반복적으로 적용합니다.
- $h^{(r)} = M_\theta(h^{(r-1)})$ (단, $r = 1, \dots, R$ )
- 최종 출력은 단계 $R$ 에서의 은닉 상태로부터 도출됩니다.
Looped Hybrid: Mamba-2 레이어와 소수의 인과적 셀프 어텐션 레이어를 모두 포함하는 백본으로 루프를 확장합니다. 이 루프는 전체 하이브리드 블록에 적용되어, 상태 공간 업데이트와 어텐션 기반 상호작용을 통해 은닉 표현의 반복적인 정교화를 허용합니다.
학습 프로토콜:
1. 사전 학습(Pre-training): 최종 루프 목적 함수(단계 $R$ 에서만 손실 계산)를 사용하여 최적화됩니다.
2. 종료 게이트 단계 1 (Exit-Gate Stage 1): 모든 재귀 단계의 출력을 사용하여 종료 깊이에 대한 토큰 수준의 분포를 학습하기 위해 엔트로피 정규화 적응형 계산 목적 함수(Ouro에서 변형)를 도입합니다.
3. 종료 게이트 단계 2 (Exit-Gate Stage 2): 백본과 헤드를 고정하고, 다음 단계가 예측 손실을 개선할 때만 계속 진행하도록 종료 게이트만을 학습시킵니다.

적응형 종료 상태 선택 (Adaptive Exit-State Selection)

본 논문은 Ouro의 2단계 종료 게이트를 Mamba에 맞게 조정했습니다. 추론 시, 임계값 $q$ 는 누적 종료 확률에 따라 각 토큰에 대한 종료 단계 $z_i$ 의 선택을 제어합니다.

의미론적 참고 사항: 저자들은 SSM에서 상태 연속성 제약 조건 때문에 $R$ 단계의 모든 상태 체인이 보존되어야 함을 명시합니다. 따라서 "종료(exit)"는 어떤 재귀 단계의 은닉 상태를 언어 모델링 헤드로 전달할지를 선택하는 것이지, 추가적인 상태 처리 메커니즘 없이는 본질적으로 실제 연산 시간(wall-clock computation)이나 실행 단계 수를 줄이지는 않습니다.

주요 기여

통제된 경험적 연구: 저자들은 비루프형(non-looped) Mamba 및 Transformer 베이스라인과 비교하여, 동일 파라미터(iso-parameter) (동일한 $N$ , 다양한 $R$ ) 및 동일 FLOPs(iso-FLOPs) (동일한 유효 깊이 $NR$) 프로토콜 하에서 Looped Mamba 및 하이브리드 모델에 대한 엄격한 비교를 제시합니다.
추론 작업 평가: 합성 추론 작업(Mano 및 $p$ -hop induction)과 다운스트림 벤치마크를 포함한 언어 모델 사전 학습을 통해 접근 방식의 유효성을 검증합니다.
적응형 종료 메커니즘: 루프형 상태 공간 모델을 위한 임계값 제어형 적응형 종료 상태 선택 메커니즘을 도입하여, 중간 깊이에서 다운스트림 성능을 향상시키는 능력을 입증합니다.

실험 결과

합성 추론 작업

Mano Task (모듈로 산술): Looped Mamba는 파라미터가 일치하는 비루프형 베이스라인보다 일관되게 우수한 성능을 보였습니다. 가장 어려운 설정( $L=24$ )에서 $4 \otimes 6$ 모델은 $4 \otimes 1$ 베이스라인의 **7.50%**와 비교하여 **99.35%**의 정확도를 달연했습니다. 동일 FLOPs 조건(유효 깊이 24)에서 $4 \otimes 6$ 모델은 $24 \otimes 1$ 비루프형 베이스라인(99.35% vs. 37.45%)을 크게 앞질렀으며, 이는 재귀적 파라미터 공유가 공유되지 않은 레이어를 늘리는 것보다 더 강력한 귀납적 편향(inductive bias)을 제공함을 시사합니다.
$p$ -hop Induction (재귀적 검색): 얕은 비루프형 Mamba 모델은 어려움을 겪었으나(예: $1 \otimes 1$ 에서 약 38-41%), 루프를 추가하면 성능이 극적으로 향상되었습니다. $1 \otimes 12$ 모델은 거의 완벽한 정확도(~99.89%)를 달성하였으며, 이는 재귀적 검색을 위해 별도의 파라미터 수보다 재귀적 깊이가 더 중요하다는 것을 나타냅니다.

언어 모델링 및 다운스트림 벤치마크

동일 파라미터 (Iso-Parameter): 고정된 파라미터 예산 하에서 재귀적 깊이는 Mamba-2 및 LLaMA-2 백본 모두의 검증 퍼플렉시티(PPL)를 개선했습니다.
동일 FLOPs (Iso-FLOPs): 더 깊은 비루프형 모델( $24 \otimes 1$ )이 일치하는 유효 깊이를 가진 루프형 모델보다 검증 퍼플렉시티 측면에서 우위를 유지했습니다. 그러나 루프형 모델은 다운스트림 벤치마크에서 경쟁력을 유지했으며, 때로는 더 깊은 비루프형 베이스라인보다 뛰어난 성능을 보였습니다.
적응형 종료 (Adaptive Exit):
- 140M 규모: 적응형 종료 게이팅은 다운스트림 성능을 개선했습니다. 140M Mamba-2 모델의 경우, Stage 2 종료 게이트는 평균 점수 38.87과 평균 선택 종료 단계 3.00을 기록하며, fixed-loop4 베이스라인(37.79)을 능가했습니다.
- 370M 규모: 최상의 다운스트림 성능은 여전히 fixed-loop4 설정 근처에 집중되어 있었습니다. 종료 게이팅 변형 모델들이 피크 루프 점수를 넘어서지 못했는데, 이는 더 큰 모델이 마지막 재귀 단계에 더 많이 의존하거나 더 나은 게이트 캘리브레이션이 필요함을 시사합니다.
- 트레이드오프: 종료 게이팅은 선택된 깊이와 성능 사이의 조절 가능한 경계를 제공하여, 모델이 고정된 최대 깊이를 강제하기보다 중간 재귀 상태를 효과적으로 활용할 수 있게 합니다.

의의 및 주장

본 논문은 재귀적 깊이가 Mamba 및 하이브리드 언어 모델을 위한 유망하고 파라미터 효율적인 스케일링 축을 제공한다고 주장합니다. 주요 발견은 다음과 같습니다:

깊이 vs. 파라미터: 많은 추론 작업은 추가적인 독립적 파라미터보다 더 큰 계산적 깊이로부터 이득을 얻습니다. Mamba 레이어를 재귀적으로 재사용하는 것은 파라미터 수를 늘리지 않고도 구성적(compositional) 및 재귀적(recursive) 계산 능력을 실질적으로 향상시킵니다.
하이브리드 생존 가능성: 루프형 깊이의 원리는 어텐션 레이어가 Mamba 레이어와 함께 공동으로 루프되는 하이브리드 Mamba-Transformer 아키텍처에도 효과적으로 적용됩니다.
적응형 선택: 현재의 SSM 구현은 전체 상태 실행을 요구하지만, 적응형 종료 상태 선택은 최적의 예측 깊이를 선택할 수 있게 하여 중간 깊이에서의 다운스트림 성능을 개선합니다(특히 140M 규모에서).

저자들은 실험이 370M 파라미터 규모의 모델까지로 제한되어 있으며, 실제 추론 시간의 연산 절감을 위해서는 재귀 단계를 완전히 건너뛰는 상태 처리 메커니즘에 대한 후속 연구가 필요하다는 점을 언급하며 자신들의 주장에 대해 신중한 태도를 유지합니다. 결과는 완전한 대규모 성능 특성화라기보다는 통제된 경험적 증거로서 제시되었습니다.