[Paper] Multiplex Thinking: 토큰별 Branch-and-Merge를 통한 추론

Source: arXiv - 2601.08808v1

Overview

이 논문은 Multiplex Thinking을 소개한다. 이는 대형 언어 모델(LLMs)이 부드럽고 확률적인 사고의 유연성과 표준 토큰 생성의 효율성을 결합한 새로운 추론 방식이다. 각 단계에서 여러 후보 토큰을 샘플링하고 이를 하나의 “multiplex” 토큰으로 병합함으로써, 모델의 어휘 지식을 유지하면서 추론 체인을 크게 단축한다. 저자들은 이 방법이 기존 Chain‑of‑Thought (CoT) 프롬프트보다 적은 토큰으로도 어려운 수학‑추론 벤치마크에서 더 강력한 성능을 보인다는 것을 입증한다.

주요 기여

• 멀티플렉스 토큰 표현: K 후보 토큰을 샘플링하고, 그 임베딩을 병합하여 결과를 단일 연속 토큰으로 취급하는 확률적 메커니즘.
• 자기 적응 행동: 모델이 자신 있을 때, 멀티플렉스 토큰은 거의 이산 토큰으로 수축하여(전통적인 CoT와 동일하게 동작); 불확실할 때는 여러 가능한 연속을 압축적으로 인코딩한다.
• 온‑폴리시 RL 최적화: 멀티플렉스 롤아웃에 대한 다루기 쉬운 확률 분포는 직접적인 강화학습 파인튜닝을 가능하게 하며, 이는 일반 이산 CoT 시퀀스로는 수행하기 어렵다.
• 실증적 향상: 여러 어려운 수학 추론 데이터셋에서 Pass@1–Pass@1024 전반에 걸쳐 강력한 이산 CoT 및 RL 베이스라인 대비 일관된 개선을 보이며, 토큰 시퀀스 길이는 더 짧다.
• 오픈소스 공개: 코드와 사전 학습된 체크포인트가 공개되어 재현성과 후속 활용을 촉진한다.

방법론

1. 샘플링 단계 – 각 추론 단계에서 모델은 소프트맥스 분포(일반 생성에 사용되는 동일한 분포)에서 K개의 후보 다음 토큰 ID를 추출합니다.
2. 임베딩 병합 – 이 K 토큰들의 임베딩을 (예: 가중 평균이나 학습된 어텐션 모듈을 통해) 하나의 멀티플렉스 임베딩으로 결합합니다.
3. 멀티플렉스 토큰 주입 – 이 임베딩을 일반 토큰처럼 트랜스포머에 다시 입력하여 토큰 수를 늘리지 않고도 모델이 추론을 계속하도록 합니다.
4. 확률 추적 – 샘플링 단계가 명시적이기 때문에 멀티플렉스 경로의 결합 확률을 분석적으로 계산할 수 있어 각 롤아웃에 대해 명확한 가능도를 제공합니다.
5. 강화 학습 파인튜닝 – 계산 가능한 가능도를 활용해 저자들은 온‑폴리시 RL(예: PPO)을 적용하여 작업별 보상(예: 수학 문제의 정답)을 직접 최대화합니다.
6. 자체 적응성 – 병합 연산은 K 샘플 토큰이 매우 집중되어(높은 신뢰도) 있을 경우 멀티플렉스 임베딩이 단일 토큰 임베딩과 거의 동일하도록 설계되었으며, 그렇지 않을 경우 여러 대안에 대한 정보를 유지합니다.

전체 파이프라인은 토큰‑임베딩 조회 주변에 작은 래퍼만 추가하면 기존 트랜스포머 API에 그대로 들어맞으며, 현재 LLM 스택에 손쉽게 연결할 수 있습니다.

Results & Findings

• Sequence length: 멀티플렉스 트래젝터리는 평균적으로 CoT 대비 약 30‑40% 짧아, 추론 지연 시간과 메모리 사용량을 감소시킵니다.
• Robustness to K: modest K (예: 3‑5)만으로도 충분한 불확실성을 포착해 성능을 향상시키며, K가 커질수록 수익이 점차 감소합니다.
• Ablation: RL 파인‑튜닝 단계를 제거하면 성능이 CoT 수준으로 다시 떨어져, 온‑정책 최적화가 멀티플렉스 롤아웃의 전체 이점을 끌어내는 데 필수적임을 확인합니다.

Practical Implications

• Faster inference for reasoning‑heavy APIs – 짧은 토큰 시퀀스는 요청당 계산 비용을 낮추어, 더 저렴하고 반응성이 뛰어난 LLM 서비스(예: 코드 자동완성, 튜터링 봇)로 직접 연결됩니다.
• Better utilization of token budgets – 모델이 최대 컨텍스트 길이(예: 온‑디바이스 추론 또는 API 토큰 제한)로 제한되는 상황에서, 멀티플렉스 사고는 더 풍부한 프롬프트나 긴 히스토리를 위한 공간을 확보합니다.
• Simplified pipeline for RL‑based alignment – 멀티플렉스 롤아웃의 확률이 계산 가능하기 때문에, 개발자는 이산 토큰 시퀀스에 사용되는 복잡한 그래디언트 추정 기법 없이도 표준 RL 알고리즘(PPO, REINFORCE)을 적용할 수 있습니다.
• Potential for multi‑modal reasoning – 동일한 멀티플렉스 개념을 비전‑언어 모델에 확장할 수 있어, 다음 언어 단계 전에 여러 시각적 가설을 병합함으로써 보다 효율적인 멀티모달 에이전트를 구현할 수 있습니다.
• Ease of integration – 이 방법은 커스텀 임베딩 레이어와 샘플링‑병합 래퍼만 필요하고, 기존 트랜스포머 가중치를 재사용할 수 있어 팀이 처음부터 재학습 없이 실험을 진행할 수 있습니다.

제한 사항 및 향후 연구

• 샘플링 오버헤드 – 단계마다 K 후보를 생성하면 순전파에 일정한 비용이 추가됩니다; 더 긴 CoT 체인보다 여전히 저렴하지만 저전력 하드웨어에서는 눈에 띌 수 있습니다.
• *K*와 병합 함수 선택 – 논문에서는 몇 가지 휴리스틱을 탐색했지만, 최적의 작업‑적응 선택 전략은 아직 해결되지 않았습니다.
• 해석 가능성 – 멀티플렉스 토큰은 명시적인 중간 추론 단계를 숨겨, 일반 CoT에 비해 디버깅이나 인간‑인‑루프 검증이 더 어려워집니다.
• 수학을 넘어선 일반화 – 실험은 산술 및 기호 추론에 초점을 맞추었으며, 멀티플렉스 사고를 개방형 QA, 코드 생성, 대화 등에 적용하려면 추가 검증이 필요합니다.
• 매우 큰 모델에 대한 확장성 – 저자들은 13B 파라미터 모델까지 테스트했으며, 70B 이상 LLM에서 이 기법이 어떻게 작동하는지는 아직 미지수입니다.

전반적으로, 멀티플렉스 사고는 소프트‑확률적 추론과 토큰 효율적인 생성의 매력적인 조합을 제공하며, 고성능·비용 인식 LLM 애플리케이션을 구축하는 개발자에게 즉각적인 이점을 약속합니다.

저자

• Yao Tang
• Li Dong
• Yaru Hao
• Qingxiu Dong
• Furu Wei
• Jiatao Gu

논문 정보

• arXiv ID: 2601.08808v1
• 카테고리: cs.CL, cs.AI, cs.LG
• 출판일: 2026년 1월 13일
• PDF: Download PDF