Tensor-Train Joint Modeling for Few-Step Discrete Diffusion

기술 요약: 소수 단계 마스크 확산(Few-Step Masked Diffusion)을 위한 텐서-트레인 결합 모델링

문제 정의
이산 확산 모델(Discrete Diffusion Models), 특히 마스크 확산 모델(Masked Diffusion Models, MDMs)은 토큰의 병렬 생성을 가능하게 함으로써 자기회귀(Autoregressive, AR) 모델의 대안으로서 순차적 이산 데이터(예: 텍스트, 분자) 생성에 유망한 대안을 제시한다. 그러나 이들의 잠재적인 소수 단계 생성 능력은 근본적인 구조적 한계인 '조건부 독립 가정'으로 인해 실현되지 못하고 있다. 현재의 MDM은 조건부 클린 분포 $p_\theta(x|x_t)$를 독립적인 주변 분포의 곱($\prod p_\theta(x_i|x_t)$)으로 모델링한다. 이 가정은 단계당 언마스킹되는 토큰 수가 많아질수록 심화되는 체계적인 "병렬화 편향(parallelization bias)"을 유발한다. 소수 단계 레지임(few-step regime)에서는 많은 토큰이 동시에 언마스킹되므로, 이 편향은 샘플 품질을 심각하게 저하시킨다.

이 편향을 해결하기 위해 결합 분포 $p_\theta(x|x_t)$를 명시적으로 모델링하는 것은 이론적으로 바람직하지만, $V^N$개의 항목을 가진 $N$차원 텐서를 표현해야 하므로 계산적으로 불가능하다(여기서 $V$는 어휘 크기, $N$은 시퀀스 길이이다). 기존 방식들은 보조 AR 모델을 사용하거나 잠재 변수를 사용하여 이를 우회하려 시도하지만, 이는 외부 추론 비용을 상속받거나 명시적인 결합 파라미터화 대신 암묵적인 상관관계에 의존한다는 한계가 있다.

방법론
저자들은 텐서 분해를 사용하여 이산 확산 모델 내부에서 결합 분포 $p_\theta(x|x_t)$를 명시적으로 파라미터화하는 프레임워크를 제안한다. 핵심 방법론은 두 가지 주요 구성 요소로 이루어진다:

1. 결합 모델링을 위한 텐서 분해:
   프레임워크는 조건부 분포를 저계수(low-rank) 텐서로 표현하며, 두 가지 특정 분해를 지원한다:

• CPD (Canonical Polyadic Decomposition): 텐서를 rank-1 텐서들의 합으로 근사한다. 이는 표준 MDM(rank-1 케이스)을 일반화하지만, 모든 토큰 위치를 대칭적으로 취급한다.
• TTD (Tensor-Train Decomposition): 텐서를 체인 형태로 연결된 코어(행렬)들의 시퀀스로 근사한다. 저자들은 TTD의 중요한 구조적 편향을 식별했다: 임의의 분할 지점에서의 TT-rank는 결합 분포의 언폴딩 행렬(unfolding matrix)의 랭크와 일치한다. 자연어 또는 분자 표기법과 같이 국소적 의존성이 지배적인 순차적 데이터의 경우, 이러한 언폴딩 랭크는 작게 유지되므로 TTD는 낮은 랭크로도 인접한 토큰 간의 의존성을 효율적으로 포착할 수 있다. 이는 Oseledets의 정리에 의해 이론적으로 뒷받침된다.

2. 반복적 주변 추론을 통한 효율적 샘플링:
   결합 분포로부터 직접 샘플링하는 것은 다루기 어렵다. 저자들은 이전 단계에서 샘플링된 위치들에 의존하여 토큰을 하나씩(또는 배치 단위로) 샘플링하는 연쇄 법칙 기반의 샘플링 절차를 제시한다.

• 일반 절차: 이 방법은 이전에 샘플링된 위치들에 조건을 부여함으로써 토큰을 반복적으로 샘플링한다. 전체 결합 분포를 평가하지 않고도 효율적으로 주변 확률을 계산하기 위해 캐싱(caching)과 병렬 접두사 합(parallel prefix sums)을 활용한다.
• 사전 결정된 스케줄: 고정된 언마스킹 스케줄의 경우, 모델은 조건화되지 않은 코어를 수축(contracting)하고 직접 수축된 코어를 예측하는 보조 헤드를 사용하여 중복 계산을 피함으로써 추론을 최적화한다.

3. 통합 및 미세 조정(Fine-Tuning):
   이 프레임워크는 경량 미세 조정을 통해 사전 훈련된 MDM에 통합되도록 설계되었다. 아키텍처 수정은 최종 출력 헤드를 텐서 분해 구조로 교체하는 것을 포함한다. 원래 모델의 주변 예측을 보존하기 위해, 가중치는 원래 헤드의 복사본(대칭성을 깨뜨리기 위한 작은 노이즈 포함)으로 초기화되며, 이를 통해 처음부터 학습하는 것에 비해 최소한의 비용으로 결합 의존성을 학습할 수 있다.

주요 기여

• 이산 확산에서의 최초 명시적 결합 모델링: 본 연구는 텐서 분해를 통해 조건부 클린 분포를 명시적으로 파라미터화하는 최초의 프레임워크를 도입하였으며, 이는 표준 MDM을 rank-one 케이스로 엄격히 일반화한다.
• 순차적 데이터를 위한 TTD의 구조적 편향: 저자들은 순차적 데이터에 대한 TTD의 적합성을 정형화했다. 그들은 TTD의 구조가 자연어 및 분자 표기법의 국소적 의존성과 일치함을 입증하여, CPD보다 우수한 성능을 보임을 보여주었다.
• 효율적인 샘플링 알고리즘: 캐싱과 병렬 접두사 합을 통해 최소한의 오버헤드로 결합 분포로부터 트래클한 샘플링을 가능하게 하는 새로운 반복적 주변 추 inference 절차를 제시한다.
• 경량 미세 조정: 사전 훈련된 모델에 이 프레임워크를 통합하는 레시피를 제공하여, 처음부터 학습하는 비용의 극히 일부만으로 상당한 소수 단계 개선을 달고한다.

실험 결과
프레임워크는 텍스트 생성(OpenWebText, LM1B) 및 분자 생성(QM9)에 대해 기본 모델인 MDLM 및 VADD를 사용하여 평가되었다.

• 텍스트 생성: TTD 기반 미세 조정은 바닐라 모델에 비해 생성 퍼플렉시티(perplexity)를 유의미하게 감소시켰으며, 특히 소수 단계 레지임에서 그러했다. 예를 들어, OpenWebText에서 TTD는 바닐라 VADD 대비 8-step 퍼플렉시티를 32.7% 감소시켰다. CPD는 베이스라인 대비 완만한 개선을 보이거나 개선이 없었다.

• 분자 생성: QM9에서 TTD 접근 방식은 베이스라인에 비해 일관되게 높은 유효성(validity), 고유성(uniqueness), 신규성(novelty) 점수를 달성했으며, 국소적 의

• 효율성: 이론적 복잡성에도 불구하고 샘플링 오버헤드는 미미했다. OpenWebText에서 TTD가 강화된 VADD는 128 타임스텝에서 원래의 VADD보다 단 1.7% 느렸다.

의의 및 주장
본 논문은 이산 확산 모델 내에서 명시적 결합 확률 모델링을 위해 텐서 분해를 성공적으로 도입한 첫 번째 사례라고 주장한다. 저자들은 저계수 파라미터화를 통해 조건부 독립 가정을 넘어섬으로써 체계적인 병렬화 편향을 완화할 수 있다고 주장한다. 저자들은 TTD가 국소적 토큰 의존성의 이론적 정렬 덕분에 순차적 데이터에 특히 적합하며, 외부 자기회귀 모델에 의존하지 않고도 고품질의 소수 단계 생성을 위한 실질적인 경로를 제공한다고 강조한다. 본 연구는 저계수 근사만으로도 유의미한 성능 향상을 위한 필요한 결합 구조를 포착하기에 충분하다는 것을 입증한다.