기술 요약: FMLM+를 통한 사후 정제 (Posterior Refinement)

문제 정의

현대의 언어 모델은 주로 자기회귀(Autoregressive, AR) 방식이며, 이는 길이 $L$의 시퀀스를 생성하기 위해 $L$번의 순차적인 순방향 패스(forward pass)를 요구합니다. 이러한 내재적인 순차적 병목 현상은 추론 속도를 늦추고 메모리 부하를 높여, 추론(reasoning) 및 강화 학습과 같은 다운스트림 애플리케이션에서의 확장성을 저해합니다.

이 문제를 해결하기 위해 비자기회귀(Non-autoregressive, NAR) 대안들이 등장했으며, 주로 두 가지 패러다임을 통해 구현되었습니다:

1. 마스크 확산 모델 (Masked Diffusion Models, MDMs): 임의의 마스킹 패턴으로 학습하여 유연한 any-order 추론을 제공합니다. 이들은 적응형 unmask-and-remask 전략을 사용할 수 있습니다. 그러나 **요인화 오류(factorization error)**로 인해 여러 토큰을 동시에 생성할 때 샘플 품질이 급격히 저하됩니다. 이로 인해 고품질 출력을 위한 적은 단계(few-step) 생성이 불가능해집니다.
2. 플로 맵 언어 모델 (Flow Map Language Models, FMLMs): 유클리드 공간 내의 가우시안 노이즈로부터 원-핫 벡터(one-hot vectors)로의 공동 수송(joint transport)을 학습하여 MDM의 요인화 오류를 해결합니다. 이들은 매우 적은 단계의 생성(심지어 단 한 번의 함수 평가만으로도)을 가능하게 합니다. 하지만 모든 토큰을 동시에 진화시키기 때문에, 생성 과정에서 신뢰도가 높은 토큰을 적응적으로 유지하는 MDM의 추론 유연성이 부족합니다.

기존의 NAR 패러다임은 인간의 초안 작성 과정(drafting process)을 모방하여, 출력의 임의의 부분 집합을 전역적으로 비판하고, 지우고, 다시 생성하여 품질을 개선하는 **반복적 정제(iterative refinement)**의 잠재력을 완전히 실현하지 못하고 있습니다. MDM은 다중 토큰 초안 작성을 위한 효율성이 부족하고, FMLM은 적응형 토큰 유지를 위한 유연성이 부족합니다.

방법론: FMLM+ 및 사후 정제 (Posterior Refinement)

저자들은 MDM과 FMLM 사이의 간극을 메우는 통합 프레임워크인 **FMLM+**와, 새로운 추론 전략인 **사후 정제(Posterior Refinement, PR)**를 소개합니다.

1. FMLM+ 아키텍처 및 학습

FMLM+는 **마스킹 스타일의 노이즈 스케줄(masking-style noise schedules)**을 갖춤으로써 Flow Map Language Models를 일반화합니다.

• 노이징 프로세스 (Noising Process): 전체 시퀀스를 균일하게 노이징하는 대신, FMLM+는 컨텍스트 위치 $C$에 있는 깨끗한 토큰 $y_C$를 보존하면서 나머지만 노이징합니다. 이는 Diffusion Transformer (DiT) 아키텍처의 시간 조건부(time conditioning)를 위치별로 수정함으로써 구현됩니다 (깨끗한 위치에 대해 $t=1$ 설정).
• 학습 목표 (Training Objective): 모델은 대각선($s=t$)을 따라 깨끗한 데이터를 재구성하고, 더 작은 점프를 교사(teacher)로 사용하여 더 큰 오프-대각선 점프($s \to t$)를 점진적으로 증류(distill)하는 점진적 자기 증류(progressive self-distillation) 목표를 사용하여 학습됩니다.
• 사전 학습된 MDM 활용: 저자들은 $s=t=0$일 때 FMLM+의 목적 함수가 MDM의 목적 함수로 축소된다는 이론적 연결 고리를 확립했습니다. 이를 통해 두 가지 가속화 전략을 사용할 수 있습니다:
  • 증류 (Distillation): 사전 학습된 MDM을 교사로 사용하여 토큰 확률을 일치시킴으로써 그래디언트 분산을 줄입니다.
  • 초기화 (Initialization): 사전 학습된 MDM 체크포인트로부터 직접 FMLM+ 가중치를 초기화합니다. 저자들은 첫 번째 적응형 레이어 정규화(adaLN) 레이어가 가우시안 노이즈 입력을 MDM이 [MASK] 토큰에 대해 생성하는 것과 동일한 표현으로 매핑함을 관찰했습니다. 이는 사전 학습된 지식을 효과적으로 보존합니다.

2. 사후 정제 (Posterior Refinement, PR)

PR은 FMLM+가 제공하는 사후(a posteriori) 신뢰도 점수를 활용하는 추론 시점의 정제 전략입니다.

• 메커니즘: 독립적인 토큰별 주변 확률(marginal)에 기반하여 신뢰도를 추정하여 다중 모드 분포에서 실패하는 MDM과 달리, FMLM+는 전체 시퀀스를 공동으로 생성합니다. 이를 통해 모델은 생성된 전체 초안(draft)을 조건으로 각 토큰의 조건부 확률을 평가할 수 있습니다.
• 루프 (The Loop):
  1. 단일 플로 맵 점프(예: $\delta_{0,1}$)를 통해 후보 시퀀스 $\hat{x}_1$을 생성합니다.
  2. 예측된 분포의 최대 확률 $p_{max}$로 정의되는 각 토큰의 사후 신뢰도를 계산합니다. 높은 $p_{max}$는 원-핫 벡터로 투영할 때 반올림 오차가 적음을 의미합니다.
  3. $p_{max}$가 임계값 이상인 토큰(또는 상위-$k$ 토큰)을 컨텍스트 세트 $C$에 **확정(Commit)**합니다.
  4. 나머지 미확정 토큰들을 **재노이징(Re-noise)**하고 과정을 반복합니다.
• 이론적 통찰: 저자들은 MDM이 알레아토리(aleatoric) 불확실성(데이터의 무작위성)과 에피스테믹(epistemic) 불확실성(모델의 확신)을 혼동한다고 주장합니다. FMLM+ 궤적은 공동 수송을 학습함으로써 진정한 에피스테믹 불확실성을 드러내며, 이를 통해 PR이 MDM이 놓치는 오류를 안정적으로 식별하고 수정할 수 있게 합니다.

주요 기여

1. 통합 프레임워크 (FMLM+): FMLM의 적은 단계 생성 효율성과 MDM의 any-order 유연성을 결합한, 마스킹 스타일 노이즈 스케줄을 갖춘 FMLM의 일반화 모델입니다. 이는 대규모 MDM 사전 학습의 발전을 활용하여 FMLM+를 확장할 수 있는 직접적인 경로를 구축합니다.
2. 사후 정제 (Posterior Refinement): 모델이 출력을 재귀적으로 스스로 수정할 수 있게 하는 새로운 사후 신뢰도 기반 전략입니다. 이는 잘못된 토큰이 일관되게 낮은 신뢰도의 생성물 집합을 형성한다는 점을 식별하여, 신뢰할 수 있는 오류 필터링 및 수정을 가능하게 합니다.
3. 광범-한 평가: FMLM+와 PR이 네 가지 벤치마크(TinyStories, OpenWebText, GSM8K(추론), Sudoku(구조적 제약))에서 MDM 및 FMLM 제품군보다 속도-품질 트레이드오프를 더 효과적으로 탐색함을 입증합니다.

결과

본 논문은 목표 성능을 달ün 데 필요한 함수 평가 횟수(NFE)를 기준으로 상당한 속도-품질 개선을 보고합니다:

• Sudoku (구조적 제약): 무조건적 Sudoku 생성에서, MDM과 사전(a priori) 신뢰도를 사용하는 FMLM+는 여러 토큰을 샘플링할 때 처참하게 실패합니다. 반면 PR을 사용하는 FMLM+는 중간 난이도 퍼즐에서 단 9 NFE만에 87.1%의 정확도를, 어려운 퍼즐에서 16 NFE만에 **71.2%**를 달성했습니다. 이는 유사한 정확도에 도달하기 위해 1024 NFE가 필요한 베이스라인 대비 32배 빠른 속도입니다.
• GSM8K (추론): PR을 사용하는 FMLM+는 GSM8K 테스트 세트에서 단 32 NFE만으로 19.0%의 정확도를 달성하여, 모든 비자기회귀 베이스라인을 능가했습니다.
• 무조건적 생성 (TinyStories/OpenWebText): FMLM+는 최대 8배 적은 NFE를 사용하여 강력한 베이스라인(AR 모델 포함)과 대등하거나 이를 능가하는 성능을 보여줍니다. 예를 들어, TinyStories에서 FMLM+는 32 NFE만으로 17.53의 생성 퍼플렉시티(generative perplexity)를 달성했는데, 이는 128 NFE를 사용한 MDLM의 18.74와 비교됩니다.
• 학습 효율성: MDM 가중치를 통한 증류 또는 초기화를 활용하면 처음부터 학습하는 것보다 학습 수렴 속도와 최종 성능이 크게 향고됩니다.

의의 및 주장

본 논문은 **사후 정제를 사용하는 FMLM+**가 비자기회귀 언어 모델링의 새로운 최첨단(SOTA)을 확립한다고 주장합니다. 그 의의는 다음과 같습니다:

• 반복적 정제의 실현: FMLM의 효율적인 다중 토큰 초안 작성 기능과 PR의 적응형 토큰 유지를 결합함으로써, NAR 모델의 이론적 이점인 '반복적 정제'를 성공적으로 구현했습니다.
• 확장성: FMLM과 MDM 사이의 간극을 메움으로써, 대규모 사전 학습된 마스크 확산 모델(예: LLaDA-8B, Dream)을 활용할 수 있는 고충실도 언어 모델링을 위한 확장 가능한 기반을 제공합니다.
• 추론 시점 스케일링 (Inference-Time Scaling): 추가적인 정제 단계를 거칠수록 샘플 품질이 향아가는, 플로 기반 모델 특유의 추론 시점 스케일링을 보여줍니다.

저자들은 실험이 소규모 및 주로 합성 데이터셋에서 수행되었으며, 처음부터 학습하는 것이 MDM 가중치로 초기화하는 것보다 최적이 아니라는 점을 언급하며 한계점에 대해 겸허한 태도를 유지했습니다. 이들은 향후 연구가 십억 파라미터 규모로의 확장과 사후 신뢰도 신호를 활용한 새로운 테스트 타임 전략 설계로 나아가야 한다고 제안합니다.