기술 요약: MUX - 멀티플렉싱된 토큰을 통한 연속적 추론 (Continuous Reasoning via Multiplexed Tokens)

1. 문제 정의

현대의 자기회귀 언어 모델은 자연어로 중간 추론 단계를 명시하는 방식(Chain-of-Thought)을 통해 복잡한 문제(수학, 코딩, 상식 등)를 해결합니다. 이 방식은 효과적이지만, 계산 측면에서는 비효율적입니다. 각 추론 단계가 단 하나의 서브워드(subword)만을 전달하기 때문에 엄격한 대역폭 병목 현상이 발생합니다. 더욱이, 많은 추론 토큰들은 최적화된 계산이라기보다 인간의 의사소통 패턴을 모방하는 중복된 정보를 담고 있습니다.

연속적 잠재 추론(Continuous latent reasoning)은 모델이 서브워드 대신 연속적인 벡터를 예측하는 대안을 제시하며, 이는 중첩(superposition)을 통해 여러 서브워드를 동시에 인코딩함으로써 이론적으로 더 높은 대역폭을 허용할 수 있습니다. 그러나 기존의 잠재 추론 방법들은 다음과 같은 큰 난관에 봉착해 있습니다:

• 전역적 방법(Global methods): 최종 정답이나 궤적의 끝점만을 감독(supervise)하는 방식은 잠재 토큰이 정보가 없는 자리 표시자(placeholder)로 변하는 "지름길 행동(shortcut behaviors)"을 초래하는 경우가 많습니다 이.
• 지역적 방법(Local methods): 각 잠재 토큰을 이산적인 추론의 특정 구간(span)에 대응시켜 감독하는 방식은 복잡한 보조 아키텍처(예: 자기회귀 디코더)를 필요로 하거나, 검색 작업에 필요한 다양한 가설의 다양성을 유지하는 데 실패하곤 합니다.

핵ales적인 과제는 이산적 추론에 충실하면서도, 잠재 붕괴(latent collapse)를 방지하고, 과도한 아키텍처 복잡성 없이 병렬 탐색을 지원할 수 있는 연속적 잠재 추론을 위한 감독 대상을 정의하는 것입니다.

2. 방법론: MUX

저자들은 MUX를 제안합니다. 이는 언어 모델이 **멀티플렉스된 토큰(multiplexed tokens)**을 사용하여 연속적인 잠재 추론을 수행하도록 학습시키는 지역적 증류(local distillation) 방법입니다.

핵심 개념

MUX는 지역적 증류 작업을 가변 길이의 범주형 신호(이산적 추론 단계의 구간)를 인코딩하는 고정 차원의 연속 신호를 학습하는 것으로 취급합니다. 통신 시스템의 코드 분할 다중화(code-division multiplexing)에서 영감을 얻은 MUX는 이산적 서브워드의 구간을 어휘 공간(vocabulary space) 내에서의 **가중 선형 중첩(weighted linear superposition)**으로 표현합니다.

기술적 구현

1. 멀티플렉스된 타겟(Multiplexed Targets): 이산적 추론 구간 $r_i = (r^1_i, \dots, r^{S_i}_i)$에 대해, 학습 타겟은 다음과 같이 구성됩니다:
   $$\text{mux}(r_i) := \sum_{j=1}^{S_i} \alpha^{(i)}_j \cdot \text{onehot}(r^j_i)$$
   여기서 $\alpha^{(i)}_j$는 정규화된 양의 위치 가중치입니다.
2. 위치 가중치 부여(Positional Weighting): 중첩이 무손실(lossless)(즉, 원래의 구간을 고유하게 복구할 수 있음)되도록 하기 위해, 저자들은 다음과 같은 특정 가중치 체계를 제안합니다:
   • 기하학적(Geometric): $w_j = \rho^{j-1}$ (지수적 감소).
   • 사인형(Sinusoidal): 구간의 끝부분 근처에서 피크를 이루는 사인 함수 기반.
   • 회전형(Rotary): 회전 위치 임베딩(rotary position embeddings)과 유사하게 평균 코사인 성분에 기반함.
   • 이론적 보장: 이러한 가중치들은 "부분 집합 합계 분리 조건(subset-sum separation condition)"을 만족하여, 서로 다른 위치 세트가 어떤 서브워드에 대해서도 고유한 총 질량(total mass)을 생성하게 함으로써, 서브워드가 여러 번 등장하더라도 모호성을 방지합니다.
3. 학습 목적 함수(Training Objective): 모델은 선형-소프트맥스 헤드와 KL 발산 손실($L_{local}$)을 사용하여 이러한 멀티플렉스된 타겟을 일치시키도록 학습됩니다. 이는 이산 추론 모델과 잠재 특징을 정렬하기 위해 표준 정답 예측 손실($L_{answer}$) 및 선택적 궤적 수준 손실($L_{global}$)과 결래됩니다.
4. 추론(Inference): 추론 과정에서 모델은 은닉 공간(hidden space)에서 잠재 토큰을 자기회귀적으로 생성합니다. 어휘 투영(vocabulary projection)은 학습 중에 감독을 위해서만 사용되며, 최종 출력 토큰을 생성하는 데는 사용되지 않습니다.

3. 주요 기여

(1) 멀티플렉스된 토큰을 통한 잠재 추론

MUX는 이산적 추론의 구간을 가중 선형 중첩으로 표현하는 단순한 지역적 증류 프레임워크를 도입합니다. 기존의 지역적 방법들과 달리, MUX는 보조 자기회귀 디코더나 압축된 캐시 타겟을 필요로 하지 않으며, 대신 어휘 공간 투영에 의존합니다.

(2) 무손실 멀티플렉싱(Lossless Multiplexing)

본 논문은 특정 위치 가중치(기하학적, 사인형, 회전형)가 무손실 멀티플렉싱을 보장한다는 이론적 증명을 제공합니다. 이는 이산적 추론 구간이 연속적 표현 안에 완전히 보존됨을 의미합니다. 저자들은 이러한 무손실성이 전역적 잠재 추론 방법에서 흔히 나타나는 실패 모드인 잠재 붕로(latent collapse)(토큰이 의미론적으로 균질화되는 현상)를 방지함을 증명합니다.

(3) 중첩을 통한 병렬 탐색

이론적으로 저자들은 멀티플렉스된 토큰이 **병렬 너비 우선 탐색(parallel BFS)**을 자연스럽게 구현할 수 있음을 보여줍니다. 토큰들이 중첩되어 있기 때문에, 단일 잠재 토큰은 여러 가설 세트(탐색 프런티어)를 동시에 표현하고 업데이트할 수 있습니다. 이를 통해 모델은 직렬적 감독 프레임워크 내에서 여러 문제 해결 경로를 병렬로 탐색할 수 있습니다.

(4) 실험적 우월성

4개의 언어 모델, 2개의 학습 코퍼스, 4개의 테스트 세트를 아우르는 32개의 평가 설정에서 MUX는 이산 및 연속 추론 모두에서 강력한 베이스라인들을 능가했습니다. 특히, MUX는 훨씬 적은 수의 토큰(약 2.4배에서 5.9배 적음)을 사용함에도 불구하고, 자신이 증류된 대상인 이산 Chain-of-Thought(SFT-CoT) 베이스라인보다 더 나은 성능을 보이는 경우가 많았습니다.

4. 실험 결과

• 수학적 추론: GSM8K-AUG, GSM8K-AUG-NL, SVAMP, GSM-Hard, MultiArith에서 MUX는 GPT-2, LLaMA 3.2 1B, 3B, 3.1 8B 모델 전반에 걸쳐 최고의 잠재 추론 성능을 달성했습니다. MUX는 종종 SFT-CoT를 능가하며, 압축된 중첩 추론이 원래의 이산 궤적보다 더 잘 일반화될 수 있음을 입증했습니다.
• 탐색 벤치마크: Minimum Nonnegative Sum (MNNS) 및 Game of 24 작업에서 MUX는 베이스라인(Coconut 및 CoT2 포함)을 크게 앞질렀으며, 이는 멀티플렉스된 토큰이 병렬 탐색 능력을 가능하게 한다는 이론적 주장을 검증합니다.
• 절제 연구(Ablation Studies):
  • 지역적 손실(Local Loss): 전역 손실($\gamma=0$)을 제거하고 멀티플렉스된 지역 손실에만 의존해도 SIM-CoT보다 우수한 결과를 얻었으며, 이는 지역적 타겟의 효능을 확인시켜 줍니다.
  • 가중치 부여(Weighting): 무손실 가중치(기하학적, 사인형, 회전형)는 손실이 있는 균등 가중치(uniform weighting)보다 일관되게 우수한 성능을 보였으나, 구조화된 데이터셋에서는 그 격차가 크지 않았습니다.
  • 청킹(Chunking): 이산적 구간을 잠재 토큰으로 무작위 청킹하는 방식이 가장 좋은 성능을 보였으며, 이는 구조화된 데이터 증강 역할을 했습니다.

5. 의의 및 주장

본 논문은 무손실 중첩을 지역적 학습 타겟으로 사용하는 것이 강력하고 효율적이며 해석 가능한 잠재 연속 추론을 달성하기 위한 충분 조건이라고 주장합니다.

• 효율성: MUX는 이산 CoT보다 적은 추론 단계로 높은 성능을 달성하여 계산적 대역폭 병목 현상을 해결합니다.
• 단순성: 이 방법은 (보조 디코더와 같은) 이전의 지역적 증류 접근 방식의 아키텍처 복잡성을 피하면서도 성능을 유지하거나 향상시킵니다.
• 해석 가능성: 프로빙(probing) 분석을 통해 학습된 잠재 토큰이 충실하고 해석 가능한 추론 내용을 인코딩함을 확인하였으며, 이는 잠재 추론이 단순히 해석 불가능한 자리 표시자라는 통념에 반합니다.
• 병렬성: 결과는 멀티플렉스된 타겟을 사용할 때 직렬적 감독으로부터 병렬 탐색 역학이 자연스럽게 창발될 수 있음을 시사하며, 이는 언어 모델의 계획 및 탐색 능력을 개선하는 새로운 경로를 제공합니다.

저자들은 단순하고 해석 가능한 지역적 타겟이 잠재 추론을 더 강력하고 학습하기 쉽게 만들 수 있으며, 이를 통해 연속 공간의 효율성과 이산적 추론의 충실성 사이의 간극을 메울 수 있다고 결론짓습니다.