기술 요약: 이산 확산 언어 모델 정제를 위한 표적의 드리프트

문제 제기

이산 확산 언어 모델 (DDLMs) 은 이산적 토큰 시퀀스를 반복적으로 제거하여 텍스트를 생성합니다. 최근 연속 생성 모델의 발전은 "드리프트 (drifting)" 방법을 도입했는데, 이는 샘플링 시 보정을 반대칭 고정점 목적 함수를 통해 학습에 흡수하는 방식입니다. 그러나 이 원리를 이산 텍스트로 이전하는 것은 간단하지 않습니다.

핵심적인 도전 과제는 이산 텍스트와 연속 드리프트 사이의 인터페이스에 있습니다. 연속 환경에서는 생성된 샘플이나 그 특징 표현을 드리프트 방향으로 미묘하게 조정하여 학습 타겟으로 사용할 수 있습니다. 그러나 DDLM 은 토큰에 대한 범주형 분포를 출력합니다. 이러한 분포를 특징 추출 전에 (argmax 또는 샘플링을 통해) 하드 토큰으로 축소하면, 결과적으로 특징 공간 손실이 모델의 로짓에 대해 미분 불가능해집니다. 따라서 표준 드리프트는 생성기를 정제하는 데 유용한 기울기를 제공할 수 없습니다. 본 논문은 다음과 같은 질문을 던집니다: 새로운 샘플러를 설계하거나 저-단계 학생 모델을 증류하지 않고도, 기존 DDLM 의 학습 목적 함수 자체를 정제하여 동일한 추론 예산에서 동일한 샘플러가 더 나은 샘플을 생성하도록 할 수 있는가?

방법론: TOKENDRIFT

저자들은 범주형 예측과 연속 특징 공간을 연결하도록 설계된 드리프트 목적 함수인 TOKENDRIFT를 수립했습니다. 이 방법은 세 가지 주요 단계를 통해 작동합니다:

1. 소프트 토큰 특징 리프트 (Soft-Token Feature Lift):
   예측을 하드 토큰으로 축소하는 대신, TOKENDRIFT 는 모델의 범주형 예측 ($p_t$) 을 소프트 토큰 특징으로 리프트합니다. 고정된 시맨틱 인코더 $\phi$와 임베딩 행렬 $E$가 주어졌을 때, 예측된 위치의 인코더 입력은 예상 임베딩 $\tilde{e}_t = p_t E$로 구성됩니다. 이는 특징 공간 손실에서 토큰 로짓까지의 미분 가능한 경로를 생성하여 시맨틱 공간에서 예측 불확실성을 보존합니다.

2. 반대칭 드리프트 추정 (Anti-Symmetric Drift Estimation):
   고정된 시맨틱 특징 공간에서, 이 방법은 인력 - 반발 구조에 기반하여 드리프트 장 $V$를 추정합니다:

   • 인력 (Attraction): 생성된 특징을 실제 데이터 ($P_{data}$) 의 근처 특징으로 끌어당깁니다.
   • 반발 (Repulsion): 생성된 특징을 현재 모델 분포 ($P_{model}$) 의 근처 특징에서 밀어냅니다.
     드리프트 벡터는 $V = b^+ - b^-$로 계산되며, 여기서 $b^+$와 $b^-$는 양 (실제) 및 음 (생성) 참조 집합의 온도 스케일된 바리센터입니다. 이 방법은 견고성을 보장하기 위해 다중 온도 체계를 사용합니다.

3. 특징 공간 고정점 목적 함수 (Feature-Space Fixed-Point Objective):
   모델은 스톱-그라디언트 타겟 $h^* = \text{sg}(h + \alpha V)$와 일치하도록 학습되며, 여기서 $h$는 현재 생성된 특징이고 $\alpha$는 드리프트 스케일입니다. 손실은 제곱 L2 거리로 정의됩니다:
   $$\mathcal{L}_{drift} = \frac{1}{2B} \sum_{i=1}^B \|h_i - h^*_i\|^2_2$$
   $h^*$가 고정되어 있으므로, 이 손실을 최소화하면 생성된 특징 $h$가 드리프트 장 $V$ 방향으로 이동합니다. 소프트 토큰 리프트를 통해 이 기울기는 로짓으로 역전파됩니다.

저자들은 또한 드리프트를 로짓 공간의 "미러 교사 (mirror teacher)" 분포로 변환하는 것과 같은 대안적 공식을 탐구했지만, 직접적인 특징 공간 매칭이 가능도와 엔트로피 사이의 가장 안정적인 균형을 제공한다는 사실을 발견했습니다.

주요 기여

• 이산 텍스트를 위한 드리프트 공식화: 이 논문은 소프트 토큰 리프트를 도입함으로써 DDLM 에 드리프트를 적용할 때 발생하는 미분 가능성 격차를 식별하고 해결했습니다. 이를 통해 특징 공간 기울기가 범주형 로짓을 업데이트할 수 있게 되었습니다.
• 정제 목적 함수: 이는 샘플러를 변경하거나 특수한 증류 파이프라인을 요구하지 않고 고정된 예산 (고정 함수 평가 횟수, NFE) 에서 생성 품질을 향상시키는 학습 측 정제 전략을 제안합니다.
• 실증적 검증: 이 방법은 두 가지 서로 다른 DDLM 백본인 마스크드 확산 (MDLM) 과 유니폼-스테이트 확산 (DUO) 에서 검증되었습니다.

실험 결과

저자들은 OpenWebText (OWT) 에서 통제된 지속 학습 실험을 수행하여 TOKENDRIFT 를 원래 체크포인트 및 동일한 추가 계산 예산을 사용하는 일반 지속 학습과 비교했습니다.

• 성능 향상: TOKENDRIFT 는 낮은 NFE 에서 생성 퍼플렉시티 (Gen.-PPL) 를 크게 향상시켰습니다.
  • MDLM(마스크드 확산) 에서 4 NFE 시, Gen.-PPL 은 기준선 대비 89% 감소했습니다 (~1942 에서 ~209 로).
  • DUO(유니폼-스테이트) 에서 4 NFE 시, Gen.-PPL 은 86% 감소했습니다 (~509 에서 ~76 로).
• 기준선 대비 비교: 향상은 미미한 증가만 보인 일반 지속 학습에 비해 상당했습니다. TOKENDRIFT 는 또한 매칭된 예산 하의 few-step 영역에서 특수 증류 방법 (예: SDTT, Di4C) 보다 우수한 성능을 보였습니다.
• 애블레이션 연구:
  • 소프트 vs 하드 토큰: 소프트 토큰 리프트를 하드 스트레이트-쓰루 추정기로 대체하면 심각한 성능 저하와 엔트로피 붕괴가 발생하여 예측 불확실성 보존이 필수적임을 확인했습니다.
  • 타겟 공간: 로짓 공간 미러 교사 변형보다 직접적인 특징 공간 L2 매칭이 더 안정적이었습니다.
  • 드리프트 균형: 반대칭 인력 - 반발 균형이 결정적이었으며, 반발만 또는 인력만 있는 변형은 치명적으로 실패했습니다.

중요성과 주장

이 논문은 TOKENDRIFT 를 기존 DDLM 을 위한 실용적인 정제 목적 함수로 제시합니다. 그 주요 중요성은 드리프트가 이산 도메인으로 효과적으로 이전되어 고정된 추론 예산 하에서 샘플 품질을 향상시킬 수 있음을 입증하는 데 있습니다.

저자들은 이 접근법이 기존 방법들에 대한 보완적인 관점을 제공한다고 주장합니다. 새로운 샘플러, 스케줄을 설계하거나 저-NFE 학생 모델을 증류하는 대신, 모델 클래스와 샘플러를 고정하고 학습 목적 함수를 통해 생성기를 직접 개선할 수 있습니다. 결과는 연속 드리프트의 인력 - 반발 구조가 이산 텍스트에 적용 가능할 뿐만 아니라, 미분 가능한 다리 (소프트 토큰 리프트) 가 유지된다면 안정적인 정제를 위해 필수적임을 시사합니다.

이 논문은 한계에 대해 겸손하게 언급하며, 효과는 고정된 시맨틱 특징 공간의 품질에 의존하며 현재 실험은 무조건적 생성에 제한됨을 지적합니다. 조건부 생성이나 더 대규모 모델로 이 방법을 확장하는 것은 향후 과제로 식별되었습니다.