💬 NLP

T $^\star$ : Progressive Block Scaling for Masked Diffusion Language Models Through Trajectory Aware Reinforcement Learning

이 논문은 AR 초기화 기반의 작은 블록 마스크 확산 언어 모델에서 트래젝토리 인식 강화 학습을 활용한 점진적 블록 크기 확장 커리큘럼인 T $^\star$ 를 제안하여, 수학 추론 성능을 유지하면서 높은 병렬 디코딩을 가능하게 한다고 설명합니다.

원저자: Hanchen Xia, Baoyou Chen, Yutang Ge, Guojiang Zhao, Siyu Zhu

게시일 2026-03-30

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Hanchen Xia, Baoyou Chen, Yutang Ge, Guojiang Zhao, Siyu Zhu

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

🎭 비유: "연극 배우의 리허설 방식"

인공지능 모델이 문제를 푸는 과정을 연극 배우가 대본을 외우고 무대에 서는 과정이라고 상상해 보세요.

1. 기존 방식 (기존의 문제점)

기존의 '마스크 확산 모델'은 한 번에 여러 줄의 대사를 동시에 외우는 방식을 시도했습니다.

작은 블록 (Block Size 4): 배우가 대사를 4 줄씩 끊어서 외웠습니다. 실수할 확률이 적고 정확했지만, 무대 위에서 한 번에 4 줄만 말하니까 시간이 많이 걸렸습니다. (비효율적)
큰 블록 (Block Size 32): 배우가 대사를 32 줄이나 한 번에 외우려고 했습니다. 이론적으로는 훨씬 빠르게 연극을 할 수 있지만, 배우가 한 번에 너무 많은 대사를 기억하려다 보니 혼란이 생기고 실수가 잦아졌습니다. (정확도 하락)

즉, "빠르게 하려면 정확도가 떨어지고, 정확하려면 느리다"는 딜레마에 빠진 것입니다.

2. T⋆의 해결책: "점진적인 리허설"

이 논문이 제안한 **T⋆**는 배우에게 **"한 번에 32 줄을 외우지 마라. 4 줄부터 시작해서, 익숙해지면 8 줄, 16 줄, 32 줄로 늘려가라"**고 가르칩니다.

단계 1 (초보자): 배우가 4 줄씩 끊어서 완벽하게 외웁니다. 이때 TRACERL이라는 '코치'가 옆에서 "이 대사는 이렇게 말하면 더 자연스러워"라고 피드백을 줍니다.
단계 2 (성장): 배우가 4 줄에 익숙해지면, 코치는 블록을 8 줄로 늘려줍니다. 배우는 이미 4 줄을 잘 외운 상태이므로, 8 줄을 외울 때도 어색함 없이 자연스럽게 적응합니다.
단계 3 (마스터): 이렇게 단계별로 블록 크기를 키워가면, 결국 32 줄을 한 번에 외워도 초보 때의 정확도를 유지하면서 초고속으로 연극을 할 수 있게 됩니다.

3. 핵심 기술: "코치의 눈 (TRACERL)"

여기서 중요한 것은 TRACERL이라는 기술입니다. 이는 단순히 정답만 알려주는 게 아니라, **배우가 대사를 외워가는 '과정' (궤적)**을 지켜보며 매 순간 "지금 이 대사를 외우는 게 맞니?"라고 질문합니다.

기존 방식 (직접 큰 블록으로 시작): 처음부터 32 줄을 주면 배우는 당황해서 망가집니다. (논문에서 말하는 '붕괴' 현상)
T⋆ 방식: 작은 블록에서 완벽하게 훈련된 뒤 큰 블록으로 넘어가므로, 배우의 뇌 (모델) 가 혼란 없이 새로운 규칙을 받아들입니다.

🚀 왜 이것이 중요한가요?

속도 vs 정확도 동시 달성: 수학 문제 같은 복잡한 추론이 필요한 일에서도, **한 번에 많은 정보를 처리 (병렬 처리)**하면서도 정답을 맞힐 확률을 높일 수 있게 되었습니다.
안정적인 성장: 갑자기 어려운 것을 시키면 망하지만, 작은 것부터 차근차근 시키면 큰 성과도 안정적으로 낼 수 있다는 것을 증명했습니다.
새로운 발견: 흥미롭게도, T⋆로 훈련된 모델은 단순히 "왼쪽에서 오른쪽으로" 대사를 읽는 기존 방식과 조금 다른, 모델 스스로가 찾아낸 새로운 생각의 순서를 사용하기도 합니다. 마치 배우가 대본을 외우는 순서를 스스로 최적화한 것과 같습니다.

💡 한 줄 요약

"T⋆는 인공지능이 복잡한 문제를 풀 때, '작은 것부터 차근차근' 연습하게 하여, '빠른 속도'와 '높은 정확도'를 모두 잡게 해주는 스마트한 학습법입니다."

이 기술이 적용되면, 앞으로 AI 가 수학 문제를 풀거나 복잡한 논리를 전개할 때 훨씬 더 빠르고 똑똑해질 것이라고 기대할 수 있습니다.

1. 문제 정의 (Problem)

배경:
최근 마스킹 확산 언어 모델 (Masked Diffusion Models, MDMs) 은 생성 속도와 병렬 처리 능력 측면에서 자기회귀 (Autoregressive, AR) 모델의 대안으로 부상하고 있습니다. 특히 SDAR(Synergistic Diffusion–Autoregression) 와 같은 모델들은 사전 훈련된 AR 모델을 기반으로 확산 모델을 초기화하여 학습합니다.

핵심 문제:
MDM 은 추론 시 '블록 단위 디코딩 (Blockwise Decoding)'을 사용하는데, 이는 블록 내 토큰들을 병렬적으로 복원 (Denoising) 하되 블록 간에는 AR 방식으로 생성하여 전역적 일관성을 유지합니다.

블록 크기 (Block Size) 의 딜레마: 블록 크기를 키우면 병렬 처리 효율이 높아져 추론 속도가 빨라지지만, 수학 추론 (Math Reasoning) 과 같은 복잡한 작업에서는 성능이 급격히 저하되는 현상이 관찰되었습니다.
학습 불안정성: 큰 블록 크기로 직접 전환하거나, 큰 블록 크기의 모델에 강화 학습 (RL) 을 적용할 경우, 모델의 성능이 붕괴 (Collapse) 되는 현상이 발생합니다. 이는 큰 블록일수록 조건부 문맥 (Conditioning Context) 이 약해지고, RL 보상이 여러 토큰에 동시에 전파되면서 확률 분포가 불안정해지기 때문입니다.

2. 제안 방법론 (Methodology)

저자들은 T⋆(T-star) 라는 새로운 학습 커리큘럼을 제안하여, 작은 블록에서 큰 블록으로 점진적으로 확장하면서도 성능 저하를 최소화하는 방법을 제시합니다.

2.1. 궤적 인식 강화 학습 (Trajectory-Aware RL, TraceRL)

기존 확산 모델의 디코딩 과정을 다단계 탈노이즈 (Denoising) 궤적으로 간주합니다.
각 탈노이즈 단계에서 최종 결정된 토큰을 '행동 (Action)'으로 정의하고, PPO(Proximal Policy Optimization) 와 유사한 목적 함수를 사용하여 정책 (Policy) 을 최적화합니다.
신호 전파: 최종 정답의 보상 (Reward) 만을 사용하는 것이 아니라, 각 탈노이즈 단계별 토큰에 대한 보상을 집계하여 단계별 이점 (Step-wise Advantage) 을 계산합니다. 이를 통해 전체 탈노이즈 궤적에 걸쳐 학습 신호가 전파되도록 합니다.

2.2. T⋆: 점진적 블록 확장 커리큘럼

T⋆ 은 TraceRL 을 기반으로 하여 블록 크기를 단계별로 늘리는 전략입니다.

초기화: 사전 훈련된 AR 기반의 작은 블록 (예: B=4) MDM 에서 시작합니다.
단계별 학습 (Stage-wise Training):
- 고정된 블록 크기 $B$ 에서 정해진 횟수 ( $K_B$ ) 만큼 TraceRL 업데이트를 수행합니다.
- 시프트된 파티션 (Shifted Partition): 표준 블록 분할뿐만 아니라, 오프셋 $\Delta = B/2$ 만큼 이동시킨 분할에서도 업데이트를 수행하여 모델이 블록 경계에 덜 의존하도록 합니다.
확장 (Expansion): 현재 블록 크기에서의 학습이 완료되면, 인접한 블록들을 병합하여 블록 크기를 2 배로 늘립니다 ( $B \leftarrow 2B$ ).
반복: 목표 블록 크기에 도달할 때까지 이 과정을 반복합니다.

이 방식은 큰 블록 크기의 모델이 약한 조건부 문맥에서 시작하는 것을 방지하고, 작은 블록에서 학습된 강력한 추론 능력을 큰 블록으로 자연스럽게 전이 (Transfer) 시킵니다.

3. 주요 기여 (Key Contributions)

T⋆ 커리큘럼 제안: MDM 에서 블록 크기를 확장할 때 발생하는 성능 저하를 해결하는 실용적이고 효과적인 학습 전략을 제시했습니다.
안정적인 확장 메커니즘: 큰 블록 크기에서 직접 RL 을 적용할 때 발생하는 '확률 밀도 이동 (Lazy Likelihood Displacement)'으로 인한 붕괴 현상을 점진적 학습을 통해 우회했습니다.
비정형 디코딩 스케줄 발견: T⋆을 통해 학습된 모델은 단순히 왼쪽에서 오른쪽으로 순차적으로 토큰을 복원하는 기존 방식 (Canonical Left-to-Right) 이 아닌, 모델이 학습한 최적의 비단조적 (Non-monotone) 탈노이즈 스케줄을 발견하여 추론 성능을 향상시킵니다.
외부 검색 없이 내부 정책 개선: 트리 검색 (Tree-of-Thoughts) 같은 외부 검색 메커니즘을 도입하지 않고, 모델 내부의 탈노이즈 정책 (Token-finalization schedule) 만을 RL 로 조정하여 추론 능력을 향상시켰습니다.

4. 실험 결과 (Results)

데이터셋 및 모델: SDAR-1.7B-Chat 및 SDAR-4B-Chat 모델을 사용하며, MATH500, GSM8K, AIME24 등 수학 추론 벤치마크에서 평가했습니다.
성능 비교:
- 직접 RL 적용 (Baseline): 큰 블록 크기 (B=16, 32) 에서 직접 TraceRL 을 적용하면 성능이 급격히 하락하거나 불안정해졌습니다.
- T⋆ 적용: 점진적 확장을 적용한 T⋆은 작은 블록 (B=4) 에서의 높은 성능을 유지하면서 블록 크기가 커져도 (B=32) 성능 저하가 거의 없거나 오히려 향상되었습니다.
- 벤치마크: MATH500, GSM8K, AIME24 에서 T⋆은 기본 모델 (Base) 과 직접 RL 적용 모델보다 일관되게 우수한 성능을 기록했습니다. (예: SDAR-1.7B 에서 B=32 시 MATH500 정확도 54.2% $\to$ T⋆ 적용 시 59.0% 향상)
효율성: 블록 크기가 커질수록 'Forward Pass 당 토큰 수 (TPF)'가 증가하여 추론 속도가 빨라졌습니다. T⋆은 이러한 병렬 처리 이점을 유지하면서 추론 정확도도 함께 향상시켰습니다.
디코딩 스케줄 분석: T⋆ 모델은 높은 LocalStrict 점수를 유지하면서도 (왼쪽에서 오른쪽 순서 유지), 완전히 결정론적이지 않은 유연한 토큰 복원 순서를 학습하여 복잡한 수학 문제 해결에 유리함을 보였습니다.

5. 의의 및 결론 (Significance)

이 논문은 확산 언어 모델이 AR 모델의 강력한 추론 능력을 유지하면서, 병렬 처리의 이점을 극대화할 수 있는 길을 제시합니다.

실용성: 대규모 언어 모델을 확산 모델로 전환할 때 발생하는 '블록 크기 확장'의 난제를 해결하여, 실제 배포 환경에서 높은 처리량 (Throughput) 과 낮은 지연 시간 (Latency) 을 요구하는 시나리오에 MDM 을 적용할 수 있는 가능성을 열었습니다.
이론적 통찰: 강화 학습을 통해 모델이 내부적으로 최적의 탈노이즈 스케줄을 학습할 수 있음을 보여주었습니다. 이는 외부 검색 도구에 의존하지 않고 모델 자체의 추론 능력을 극대화하는 새로운 방향성을 제시합니다.
한계 및 향후 과제: 현재 연구는 매우 큰 블록 (B=64 이상) 까지는 확장하지 못했으며, 추론 엔진의 불안정성 문제가 남아있습니다. 또한, T⋆이 성능 저하를 완화하지만 완전히 제거하지는 못한다는 점은 향후 '고품질 콜드스타트 (Cold-start)' 단계 개발의 필요성을 시사합니다.

요약하자면, **T⋆**은 확산 언어 모델이 수학 추론과 같은 복잡한 작업에서 큰 블록 크기를 사용하더라도 성능을 잃지 않도록 하는 점진적 학습 커리큘럼을 성공적으로 입증한 연구입니다.