기술 요약: 적은 단계 생성을 위한 잠재 커널 이산 플로우 맵 (Latent-Kernel Discrete Flow Maps)

문제 정의

이산 확산(Discrete diffusion) 및 플로우 매칭(Flow-matching) 모델은 여러 단계를 거쳐 노이즈를 제거함으로써 텍스트를 생성합니다. 계산 효율성을 유지하기 위해, 이러한 모델들은 일반적으로 현재 시퀀스 상태가 주어졌을 때 각 토큰 위치를 독립적으로 결정하며 전이 확률을 인수분해(factorize)합니다. 이러한 인수분해는 개별 단계의 비용을 저렴하게 만들지만, **적은 단계 생성(few-step generation)**에 있어서는 근본적인 병목 현상을 야기합니다.

대상 시퀀스가 위치 간의 강한 상관관계(예: 주어-동사 일치 또는 괄호 짝 맞추기)를 요구할 때, 독립적인 업데이트는 토큰들을 별개로 확정 짓게 됩니다. 이는 국소적으로는 일관되지만 전역적으로는 불일치하는 조합(예: "dog barks"와 "dogs bark" 사이의 혼선)을 초래할 수 있습니다. 이러한 불일치를 수정하려면 추가적인 디노이징 단계가 많이 필요하며, 이는 적은 단계 생성의 효율성 이점을 상쇄합니다.

기존의 가속화 솔루션은 두 가지 범주로 나뉩니다:

1. 우도 모델 (Likelihood models) (예: MDLM, SEDD): 강력한 인수분해 디노이저를 처음부터 학습시키지만, 상관관계를 해결하기 위해 추론 시 많은 단계가 필요합니다.
2. 적은 단계 가속기 (Few-step accelerators) (예: Di4C, SDTT, ReDi): 학습된 "교사(teacher)" 모델을 몇 단계의 "학생(student)" 모델로 증류(distill)하거나 교정(rectify)합니다. 그러나 이러한 방법들은 교사 모델의 품질 상한선을 그대로 물려받으며, 모방하는 모델의 성능을 뛰어넘을 수 없습니다.

핵심 과제는 모델이 교사에 의존하지 않고 어떻게 상관된 단계(correlated steps)를 본질적으로 표현할 수 있는지, 즉 인수분해 가정을 근본적으로 깨뜨리면서도 낮은 계산 비용을 유지할 수 있는지 여부입니다.

방법론: 잠재 커널 이산 플로 맵 (Latent-Kernel Discrete Flow Maps, LKF)

저자들은 단일 공유 잠재 변수에 의해 결합된 인수분해된 성분들의 혼합(mixture)을 도입하여 단계 내 상관관계를 도입하는, 처음부터 학습하는 플로우 맵 모델인 LKF를 제안합니다.

모델 아키텍처

유한 시간 전이 커널 $\Psi_{t \to s}(x_s | x_t)$는 다음과 같이 $M$개의 인수분해된 성분의 혼합으로 정의됩니다:
$$\Psi_{t \to s}(x_s | x_t) = \sum_{k=1}^{M} w_k(x_t, t, s) \prod_{i=1}^{L} P_i(x_s^i | x_t, k, t, s)$$
여기서:

• $k \in [M]$는 시퀀스당 한 번 샘플링되는 공유 잠재 성분 인덱스입니다.
• $w_k$는 성분 $k$의 확률을 예측하는 라우터 가중치입니다.
• $P_i$는 잠재 변수 $k$에 조건화된 위치 $i$의 인수분해된 토큰 분포입니다.

핵심 메커니ism:

1. 공유 잠재 변수: 디노이징 궤적의 시작 시점(또는 각 단계)에서, 라우터 $w(\cdot | x_t)$로부터 단일 잠재 변수 $k$가 추출됩니다.
2. 조건부 독립성: $k$가 주어지면 토큰들은 여전히 위치 전반에 걸쳐 독립적으로 샘플링됩니다. 그러나 $k$가 전체 시퀀스에 걸쳐 공유되기 때문에, $k$에 대한 주변화(marginalization)는 위치 간의 상관관계를 유도합니다.
3. 효율성: 모델은 지연 잠재 주입(late latent injection) 아키텍처를 사용합니다. 공유 트랜스포머 트렁크(trunk)가 입력을 처리하여 잠재 변수와 무관한 표현을 생성합니다. 오직 마지막 $L_k$개의 블록만이 $k$에 조건화되어 $M$개의 성분을 위해 복제됩니다. 이를 통해 주변화의 계산 비용을 낮게 유지합니다 ($O(M \cdot \text{depth})$가 아닌 $O(M \cdot L_k)$).

학습 및 추론

• 학습: 모델은 잠재 변수 $k$를 정확하게 주변화하여 학습됩니다. 목적 함수는 혼합의 로그 우도입니다:
  $$\log \Psi = \log \sum_{k} \exp(\log w_k + \ell_k)$$
  여기서 $\ell_k$는 대상 토큰에 대해 성분 $k$가 할당한 로그 확률입니다. 이는 라우터가 잠재 변수가 시퀀스 수준의 상관관계를 포착하는 혼합을 학습하도록 강제합니다. 라우터 정규화는 혼합 성분들이 활용되도록(단일 성분으로 붕괴되는 것을 방지) 하고, 잠재 변수가 충분한 정보를 전달하도록 보장합니다.
• 추론 (Best-of-M): 혼합의 이질성을 활용하기 위해, 저자들은 Best-of-M 디코딩 전략을 채택합니다. $M$개의 병렬 롤아웃이 생성되며, 각 롤아웃은 서로 다른 성분 $k$를 고정합니다. 최종 출력은 모델 자체의 시퀀스 수준 혼합 음의 ELBO(negative ELBO, NELBO)를 기반으로 선택됩니다. 이 선택 과정에는 외부 보상 모델이나 교사가 필요하지 않습니다.

주요 기여

1. 본질적 상관 단계 (Native Correlated Steps): 단일 공유 잠재 변수에 의해 결합된 $M$개의 인수분해된 성분들의 혼합을 통해 상관된 단계를 본질적으로 표현하는 플로우 맵 커널인 LKF를 도입하여, 교사 모델의 필요성을 제거했습니다.
2. 이론적 분석: 랭크-$M$ 단계가 가질 수 있는 총 상관관계(total correlation, TC)는 $(L-1)H(k|xt)$에 의해 제한된다는 정확한 정보 분해를 제시합니다. 저자들은 모델의 근사 오차가 대상 분포의 포착되지 않은 총 상관관계에 의해 하한이 됨을 증명합니다.
3. 합성 데이터 검증: 합성 데이터셋(hidden-agreement 및 parity)에 대한 실험을 통해, LKF가 표현 가능한 상관관계에 대해 $M$이 증가함에 따라 우도(likelihood)를 단조적으로 개선하며 우도를 향상시킨다는 것을 보여주었습니다. 또한, 지수적 혼합 랭크(exponential mixture rank)를 요구하는 상관관계(parity)를 포착하는 데 실패함을 보여줌으로써 이론적 난이도 결과와 일치함을 입증했습니다.
4. 실증적 성능: LM1B 및 WikiText-103 벤치마크에서 LKF는 최첨단(SOTA) 적은 단계 생성 Perplexity를 달s성했습니다. $M=8$일 때, LKF는 우도 베이스라인(MDLM, SEDD)보다 생성 Perplexity(gen-PPL)를 2.1배에서 3.3배 개선했으며, 교사 없이도 증류/교정 가속기(Di4C, ReDi)를 능가했습니다.

실험 결과

• 벤치마크: LM1B 및 WikiText-103에서 평가되었습니다.
• 성능:
  • LM1B: $M=8$ LKF 모델은 32회의 함수 평가(NFE)에서 105의 생성 Perplexity(gen-PPL)를 달성했습니다. 이는 MDLM의 304, 가장 우수한 교정 가속기인 ReDi의 117과 비교됩니다.
  • WikiText-103: $M=8$ 모델은 75의 gen-PPL에 도달하여, MDLM(228) 및 SEDD(247)를 크게 앞질렀습니다.
• 다양성: 엔트로피 붕괴(낮은 다양성의 텍스트 생성)를 겪기 쉬운 교정 가속기와 달리, LKF는 우도를 개선하면서도 높은 샘플 엔트로피를 유지합니다.
• 절제 연구 (Ablations):
  • 이 이득은 단순히 리랭킹(reranking) 때문이 아니라 혼합 구조 자체 덕분입니다. $M$개의 성분을 단일 성분의 $M$개 노이즈 드로우(noise draws)로 대체하면 성능이 크게 떨어집니다 (예: $M=8$에서 PPL 35% 상승).
  • Best-of-M 선택이 핵심입니다. 단일 결정 정책(commit-M)은 베이스라인과 유사한 결과를 보여, 혼합 구조가 선택을 활용할 수 있는 후보 다양성을 제공함을 확인시켜 줍니다.
  • 라우터 효율성: 라우터는 품질 신호라기보다 혼합 가중치 역할을 합니다. 선택은 라우터의 초기 예측보다는 모델 자신의 우도(running evidence)를 기반으로 할 때 가장 효과적입니다.

의의 및 주장

이 논문은 LKF가 "교사 의존적" 가속에서 "처음부터 학습하는" 적은 단계 생성으로의 패러다임 전환을 나타낸다고 주장합니다. 전이 커널을 저차원 혼합(low-rank mixture)으로 모델링함으로써, LKF는 인수분해 모델과 동일한 샘플링 복잡도로 상관된 완성(예: 주어-동사 일치)에 질량을 배치할 수 있게 합니다.

저자들은 다음을 강조합니다:

• 이 방법은 교사의 품질 상한선을 물려받지 않습니다.
• 이득은 혼합 크기 $M$에 따라 확장됩니다.
• 접근 방식은 이론적으로 근거가 있으며, 상관 용량 및 근사 오차에 대한 경계치를 가집니다.
• 표준 언어 모델링 벤치마크에서 베이스라인 방법들 중 최고의 적은 단계 생성 Perplexity를 달성하여, 증류나 교정의 오버헤드 없이 효율적이고 고품질의 이산 생성을 위한 실행 가능한 경로를 제시합니다.