기술 요약: 세트 디퓨전 (Set Diffusion)

문제 정의

이산 디퓨전 모델(Discrete diffusion models)은 자기회귀(Autoregressive, AR) 언어 모델의 경쟁력 있는 대안으로 부상하며, 더 빠른 생성 속도와 개선된 제어 능력을 제공한다. 그러나 기존의 디퓨전 방식은 다음과 같은 중대한 한계에 직면해 있다:

1. 고정 길이 및 양방향 제약: 대부분의 이산 디퓨전 모델은 완전한 양방향 컨텍스트를 요구하는 아키텍처에 의존하므로, 키-값(Key-Value, KV) 캐싱 사용을 방지하고 생성을 고정된 길이에 제한한다.
2. 블록 디퓨전(Block Diffusion)의 경직성: 블록 디퓨전(BD3LM)은 고정된 크기의 시퀀스 블록을 왼쪽에서 오른쪽 방향으로 생성함으로써 격차를 부분적으로 해소하지만, 유연성이 부족하다. 이는 임의 위치 삽입(infilling)이나 블록 내 가변 길이 생성을 지원할 수 없으며, KV 캐시 업데이트가 전체 블록이 디코딩될 때까지 지연된다.
3. 트레이드오프의 한계: 현재의 방법들은 주로 고정된 블록 크기를 변경함으로써 자기회귀와 디퓨전 사이를 보간(interpolate)하는데, 이는 디코딩의 유연성과 병렬성을 제한한다.

핵리적인 과제는 단순히 블록 크기가 아니라, 토큰 집합이 생성되는 **순서(order)**를 변화시킴으로써 자기회귀와 디퓨전 사이를 보간하여, 높은 샘플 품질을 유지하면서도 유연한 길이의 생성, KV 캐싱, 그리고 임의의 순서 디코딩을 가능하게 하는 언어 모델을 개발하는 것이다.

방법론

저자들은 고정된 블록 대신 유연한 위치와 유연한 길이의 토큰 집합에 대해 생성 과정을 분해하는 세트 디퓨전(Set Diffusion) 프레임워크를 소개한다.

1. 확률 모델

세트 디퓨전은 토큰 위치의 부분 집합인 $\sigma = (\sigma_1, \dots, \sigma_N)$에 대한 자기회귀적 확률 분포를 정의한다. 여기서 각 $\sigma_n$은 토큰 위치들의 부분 집합이다. 가능도(Likelihood)는 다음과 같이 분해된다:
$$p_\theta(x) = \sum_{\sigma} \prod_{n=1}^N \pi(\sigma_n | x_{<\sigma_n}) \cdot p_\theta(x_{\sigma_n} | x_{<\sigma_n})$$
여기서 $\pi$는 다음 위치 집합을 선택하는 순서 정책(order policy)이며, $p_\theta$는 디퓨전을 사용하여 해당 집합 내 토큰의 조건부 가능도를 모델링한다. 모델은 이러한 집합들에 대한 디퓨전 손실을 합산하는 음의 증거 하한(Negative Evidence Lower Bound, NELBO)을 최적화한다.

2. 세트-인과적 아키텍처(Set-Causal Architecture) 및 KV 캐싱

핵심적인 아키텍처 혁신은 **세트-인과적 트랜스포머(set-causal transformer)**이다. 표준 양방향 디퓨전 모델과 달리, 이 아키텍처는 매 추론 단계마다 KV 캐싱을 지원한다.

• 메커니즘: 토큰 집합 $x_{\sigma_n}$이 생성되면, 그 키(key)와 값(value) 벡터는 캐싱되어 이후의 예측을 위한 조건으로 재사용된다.
• 학습 효율성: 저자들은 특수화된 어텐션 마스크를 사용하여 단일 순방향 패스(forward pass) 내에서 깨끗한 컨텍스트와 노이즈가 섞인 쿼리를 동시에 처리하는 벡터화된 학습 알고리즘을 제안한다. 이 마스크는 특정 집합 내의 노이즈 섞인 토큰이 동일한 집합 내의 다른 노이즈 섞인 토큰들과, 그리고 이전 집합들의 모든 깨끗한 토큰들만을 어텐션하도록 보장한다.

3. 토큰 순서 보간 (Interpolating Token Orderings)

샘플 품질(경계의 타이트함)과 추론 효율성(병렬성) 사이의 트레이드오프를 제어하기 위해, 저자들은 **위치 의존적 드러냄 시간 스케줄(position-dependent reveal-time schedules)**을 도입한다.

• 위치-오프셋 스케줄(Position-Offset Schedules): 각 토큰 위치 $\ell$은 누적 분포 함수(CDF) $\alpha^\ell_\tau$로 정의된 분포에서 추출된 무작위 드러냄 시간 $R_\ell$을 할당받는다.
• 슬라이딩 윈도우 인스턴스화 (SW-SetDLM): 저자들은 토큰들이 활성 생성 구간 $[a_\ell, a_\ell + w]$를 갖는 특정 인스턴스화를 제안한다. 윈도우 크기 $w$는 "왼쪽에서 오른쪽" 편향을 제어한다:
  • 작은 $w$ (1/L에 근접): 순서를 왼쪽에서 오른쪽 생성에 집중시켜 AR 동작을 회복하고 ELBO를 타이트하게 만든다.
  • 큰 $w$ (1에 근접): 순서에 구애받지 않는 생성을 허용하여 표준 마스크 디퓨전(masked diffusion)을 회복한다.
  • 중간 값: 슬라이딩 윈도우 디코딩을 가능하게 하여, 모델이 임의의 위치에서 유연한 길이의 토큰 집합을 예측하고 인필링(infilling)과 같은 작업을 지원할 수 있게 한다.

4. 추론

추론 과정에서 모델은 현재의 순서 시간과 위치-오프셋 스케줄을 기반으로 후보 위치들을 샘да링한다. 그 후 KV 캐시를 사용하여 이러한 위치들을 병렬적으로 디노이징(denoise)한다. 수용된 토큰들은 확정(commit)되며, 캐시는 즉시 업데이트되어 가변 길이 생성과 효율적인 디코딩을 가능하게 한다.

주요 기여

1. 세트 디퓨전 프레임워크: 유연한 위치와 유연한 길이의 토큰 집합에 대해 생성을 분해하는 새로운 클래스의 언어 모델로서, 자기회귀, 순서 불가지론적(order-agnostic) 디퓨전, 그리고 블록 디퓨전을 일반화한다.
2. 세트-인과적 아키텍처: 매 추론 단계마다 KV 캐시 업데이트를 지원하여 효율적인 가변 길이 생성을 가능하게 하는 트랜스포머 아키텍처이다.
3. 순서 보간: 고정된 블록 크기가 아닌 위치 의전적 순서 스케줄(특히 슬라이딩 윈도우 스케줄)을 튜닝함으로써 AR과 디퓨전 사이를 보간하는 원칙적인 방법을 제공한다.
4. 최첨단 성능(SOTA): 세트 디퓨전이 수학적 추론, 요약, 무조건적 생성(unconditional generation)에서 기존 디퓨전 모델보다 우수한 속도-품질 트레이드오프를 달성함을 입증하였으며, 인필링 작업에서 블록 디퓨전을 크게 앞선다는 것을 보여주었다.

실험 결과

저자들은 여러 벤치마크에서 세트 디퓨전(특히 SW-SetDLM)을 평가하였다:

• 수학적 추론 (GSM8K): 세트 디퓨전은 블록 디퓨전보다 엄격하게 더 나은 속도-정확도 파레토 프런티어(Pareto frontier)를 달성한다. 예를 들어, 디코딩 윈도우 크기 $S \le 8$일 때, 60.42 tokens/sec 이상의 처리량으로 66.41%의 pass@1 정확도를 달성하며, 이는 $S=8$일 때 55.39 tokens/sec에서 63.53%의 정확도를 보이는 블록 디퓨전을 능가한다.
• 인필링 (ROCStories): 세트 디퓨전은 인필링 작업(중간 문장 1개 또는 3개 채우기)에서 블록 디퓨전을 실질적으로 압도하며, 약 25% 더 빠른 디코딩 속도와 함께 더 높은 ROUGE 점수를 달升한다. 블록 디퓨전은 경직된 왼쪽-오른쪽 블록 구조로 인해 이 작업에서 어려움을 겪는다.
• 요약 (CNN/DailyMail): 세트 디퓨전은 디퓨전 베이스라인들과 경쟁력 있는 ROUGE 점수를 달성하며, 블록 디퓨전보다 최대 10% 더 빠른 디코딩 속도를 보인다. 또한 저자들이 언급했듯이 과적합 경향이 있는 AR 베이스라인보다도 우수한 성능을 보인다.
• 무조건적 생성 및 가능도 (OpenWebText 및 LM1B): OpenWebText에서 세트 디퓨전은 블록 디퓨전보다 높은 MAUVE 점수를 달성하면서도 더 빠르게 디코딩한다. LM1B에서는 기존 디퓨전 모델보다 높은 처리량을 유지하면서도 가장 좋은 퍼플렉시티(27.83 PPL for $S \le 8$)를 달성한다.

의의 및 주장

본 논문은 세트 디퓨전이 자기회귀와 디퓨전 사이의 보간을 경직된 블록 구조로부터 분리함으로써 중요한 진전을 이루었다고 주장한다. 토큰 순서를 통해 보간함으로써, 이 프레임워크는 다음을 제공한다:

• 더 큰 유연성: 임의 위치 삽입 및 가변 길이 생성을 지원하며, 이는 인필링과 같은 작업에 매우 중요하다.
• 효율성: 매 단계마다 KV 캐시를 업데이트할 수 있는 능력은 기존에 디퓨전 모델에서 불가능했던 기능으로, 더 빠른 추론을 가능하게 한다.
• 성능: 기존 디퓨전 언어 모델에 비해 우수한 샘플 품질-추론 속도 트레이드오프를 제공한다.

저자들은 세트 디퓨전을 AR, 순서 불가지론적 디퓨전, 블록 디퓨전을 특수 사례로 회복하는 통합 프레임워크로 포지셔닝하며, 추론 작업의 특정 요구 사항에 따라 병렬성과 샘플 품질 사이의 균형을 맞출 수 있는 튜너블한 메커니즘을 제공한다.