기술 요약: Unifusion

문제 정의

사전 학습된 자기회귀(Autoregressive, AR) 언어 모델을 이산 확산(Discrete Diffusion) 모델로 적응시키는 기존 방식들은 주로 **흡수 마스크 오염(Absorbing Mask Corruption)**에 의존합니다. 이러한 선택은 AR 학습과의 자연스러운 연결성(AR 디코딩은 특정 토큰 노출 순서의 특수한 형태임)을 제공하지만, 중대한 한계를 가집니다. 즉, 일단 토큰이 노출되면 샘플링 과정 중에 이를 수정할 수 없다는 점입니다. 이는 샘플링 예산(Budget)이 늘어남에 따라 반복, 템플릿 붕괴, 또는 무의미한 파편 생성과 같은 생성 문제를 야기하곤 합니다.

반면, **균등 노이즈 확산(Uniform-noise Diffusion)**은 토큰을 특수 마스크가 아닌 어휘 항목(Vocabulary items)으로 대체하여, 모든 위치가 샘플링 내내 편집 가능한 상태를 유지하게 합니다. 이는 더 높은 어휘적 다양성과 제어 가능성을 제공할 잠재력이 있습니다. 그러나 기존의 이산 확산 모델들(SEDD, MDLM/GIDD, M2S, Neural CTMC)은 서로 다른 예측 파라미터화(예: concrete scores, clean tokens, posterior means, exit rates)와 별개의 손실 함수를 사용하기 때문에, 사전 학습된 AR 체크포인트를 균등 확산으로 직접 적응시키는 것은 매우 어려웠습니다. 이러한 차이점은 각 방법론이 근본적으로 다른 역과정(Reverse process)을 학습하는 것인지, 아니면 단지 서로 다른 출력을 통해 동일한 과정을 설명하고 있는 것인지를 불분명하게 만들며, AR 출력을 균등 커널(Uniform kernels)로 직접 재사용하는 것을 방해합니다.

방법론

1. 역율 공간(Reverse-Rate Space)에서의 이산 확산 통합

저자들은 네 가지 주요 이산 확산 프레임워크인 SEDD, MDLM/GIDD, M2S, Neural CTMC 사이의 이론적 연결 고리를 확립합니다.

• 통합 목적 함수: 저자들은 이 네 가지 방법의 조건부 손실이 모두 **모델 역율(Model reverse rates)**에 대한 단일 일반화된 쿨백-라이블러(KL) 목적 함수로 표현될 수 있음을 입증합니다. 각 방법의 구체적인 출력(scores, tokens, means, rates)을 유도된 비대각 역율(Off-diagonal reverse rates)로 매핑함으로써, 근본적인 최적화 대상이 동일함을 보여줍니다.
• 변환 정리: 논문은 **클린 토큰 예측($x_0$)**을 SEDD, M2S, Neural CTMC에서 요구하는 특정 파라미터화로 변환하는 명시적인 수학적 변환식을 도출합니다. 이를 통해 단일 $x_0$ 예측이 이 모든 파라미터화를 구동할 수 있는 공유 인터페이스를 생성합니다.
• 커널 불가지론(Kernel Agnosticism): 클린 토큰 분포 $x_\theta(x_0 | x_t, t)$는 오염 커널과 관계없이 "오염되기 전 이 위치에 어떤 어휘 토큰이 있었는가?"라는 동일한 질문에 답합니다. 일단 예측되면, 알려진 순방향 커널(흡수형 또는 균등형)이 이 예측을 필요한 역율로 변환합니다. 이는 AR 초기화와 균등 확산 사이의 간극을 메워줍니다.

2. Unifusion 접근 방식

이러한 이론적 연결을 바탕으로, 저자들은 중간 단계인 마스크 확산 단계를 거치지 않고 사전 학습된 GPT-2 체크포인트를 균등 노이즈 확산으로 직접 적응시키기 위한 간단한 지속적 사전 학습(Continual Pre-training) 전략인 Unifusion을 제안합니다.

• 초기화: 이 방법은 AR 모델의 내재된 구조를 활용합니다. 위치 $i$에 대한 AR 체크포인트의 다음 토큰 분포는 $x_{<i}$를 기준으로 한 칸 밀려(Shifted) 위치 $i$에서의 클린 토큰 $x_0$에 대한 초기 예측값으로 사용됩니다.
• 학습 역학:
  • 어텐션 어닐링(Attention Annealing): 모델은 인과적 어텐션(Causal attention, AR 학습에 사용됨)에서 가시적인 오염된 토큰들에 대한 완전한 양방향 어텐션(Bidirectional attention)으로 전환됩니다. 이를 통해 최적화 과정 동안 전체 컨텍스트를 점진적으로 개방할 수 있습니다.
  • 시간 조건화(Time Conditioning): AR 체크포인트에는 시간 조건화가 결-여되어 있으므로, 모델은 다른 가중치들은 AR 초기화를 유지하면서 미세 조정(Fine-tuning) 중에 시간 의존적 파라미터를 학습합니다.
  • 오염 스케줄: 모델은 선형 균등 오염 스케줄($\alpha_t = 1-t, \beta_t = t$)과 균등 사전 분포(Uniform prior)를 사용하여 학습됩니다.

핵심 기여

1. 이론적 통합: 본 논문은 역율 공간에서 SEDD, MDLM/GIDG, M2S, Neural CTMC의 조건부 손실을 통합하여, 동일한 순방향 CTMC 및 노이즈 스케줄 하에서 이들이 동등함을 증명했습니다. 또한 클린 토큰 예측과 네 가지 파라미터화 사이의 명시적인 변환 공식을 제공합니다.
2. 직접적인 AR-to-Uniform 적응: 클린 토큰 $x_0$ 예측을 가교로 사용함으로써, 저자들은 중간 마스크 확산 체크포인트를 거치지 않고 사전 학습된 AR 모델을 균등 확산 커널로 직접 적응시킬 수 있게 했습니다.
3. 실증적 검증: 연구는 역율 공간의 동등성과 유한 예산 최적화(Finite-budget optimization)를 분리하여, 이론적 동등성은 존재하지만 실제 학습 궤적은 다르다는 것을 보여줍니다. Unifusion은 직접적인 AR-to-uniform 학습이 두 단계 마스크 경로와 유사한 성능을 달성하면서도 더 우수한 트레이드오프를 보임을 입증합니다.

결과

저자들은 GPT-2에서 파생된 124M(Unifusion-S) 및 355M(Unifusion-M) 파라미터 모델에 대해 Unifusion을 평가했습니다.

• 생성 퍼플렉시티(GenPPL) 대 엔트로피: Unifusion은 샘플링 예산을 16에서 256 스텝으로 늘림에 따라 GenPPL과 유니그램 엔트로피 사이의 트레이드오프를 꾸준히 개선합니다.
  • 256 스텝에서 Unifusion-S는 엔트로피 5.2626과 함께 97.783의 GenPPL을 달성했습니다.
  • 256 스텝에서 Unifusion-M은 엔트로피 5.6669와 함께 71.516의 GenPPL을 달성했습니다.
  • 동일한 규모의 평가된 모델 중 Unifusion만큼 두 지표를 동시에 우수하게 수행한 모델은 없었습니다. 예를 들어, 일부 베이스라인은 더 낮은 GenPPL을 달성했지만, 엔트로피가 현저히 낮아(반복 또는 붕괴를 의미) 문제가 되었습니다. 반면 Unifusion은 높은 어휘적 다양성을 유지했습니다.
• 다운스트림 정확도: Unifusion 모델은 WinoGrande, SIQA, BBH에서 비교된 확산 모델 중 가장 높은 제로샷 정확도를 기록했으며, 이는 GPT-2 초기 모델 및 DiffuGPT, SEDD와 같은 다른 확산 베이스라인보다 뛰어난 성능입니다.
• 샘플링 효율성: 64-스텝의 Unifusion-M 지점은 이미 256-스텝의 DiffGPT-M 지점을 파레토 우위(Pareto-dominate)에 놓았으며, 이는 샘플링 예산 측면에서 높은 효율성을 시사합니다.

의의 및 주장

본 논문은 Unifusion이 AR 체크포인트로부터 직접 균등 노이즈 확산 모델에 도달하는 방법이라는 "열린 질문"을 성공적으로 해결했다고 주장합니다. 저자들은 다음과 같은 점을 입증했습니다:

1. 균등 확산은 중간 마스크 단계 없이도 AR 초기화를 효과적으로 상속받을 수 있습니다.
2. 직접적인 적응은 기존 방식에 비해 유창성(낮은 GenPPL)과 다양성(높은 엔트로피) 사이에서 더 우수한 균형을 제공합니다.
3. 서로 다른 이산 확산 손실의 이론적 동등성은 역율 공간에서 일관된 최적화 목표를 유지하면서도 유연한 구현 선택을 가능하게 합니다.

이 연구는 오염 커널의 선택(마스크 vs 균등)이 생성 품질에 영향을 미치는 결정적인 설계 요소이며, AR 모델로부터 적절히 적응된 균등 커널이 마스크 확산에서 흔히 발생하는 반복 문제를 피하고 더 견고한 생성 과정을 제공한다는 것을 시사합니다.