🤖 machine learning

Accelerating Discrete Diffusion Models with Parallel-In-Time Sampling

이 논문은 합성, 이미지 및 텍스트 작업 전반에서 생성 품질을 유지하면서 시간 복잡도와 실행 시간을 크게 줄이기 위해 $\tau$ -leaping 알고리즘의 연속 시간 확률 적분 형태와 피카르 반복법(Picard iteration)을 활용하는 이산 확산 모델을 위한 병렬 시간 샘플링 방법을 소개한다.

원저자: Yu Yao, Huanjian Zhou, Andi Han, Wei Huang, Masashi Sugiyama

게시일 2026-07-02

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Yu Yao, Huanjian Zhou, Andi Han, Wei Huang, Masashi Sugiyama

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

당신이 파쇄된 문서를 재구성하려고 노력하고 있다고 상상해 보세요. 다만 종이가 아니라, 문장이나 이미지가 거대한 물음표(즉, "마스크")로 서서히 변해버린 상태를 복구하는 것입니다. 이것이 바로 **이산 확산 모델(Discrete Diffusion Models)**이 작동하는 방식입니다. 이 모델들은 깨끗한 그림이나 텍스트에서 시작하여, 그것을 노이즈(마스크)로 바꾼 뒤, 컴퓨터가 그 과정을 역으로 수행하여 원래의 것을 재현하는 법을 배웁니다.

문제는 무엇일까요? 현재 컴퓨터가 이 작업을 수행하는 방식은 마치 한 사람이 문서를 한 단어씩, 엄격한 순서에 따라 재구성하는 것과 같습니다. 첫 번째 단어를 추측하고, 그다음 두 번째, 그다음 세 번째를 추측해야 합니다. 설령 수천 개의 코어를 가진 초고속 컴퓨터(최신 GPU와 같은)를 가지고 있더라도, 이 방식은 한 단계가 끝나야만 다음 단계를 시작할 수 있도록 컴퓨터를 강제합니다. 이는 마치 다음 주자가 움직이기 전에 반드시 바통을 완벽하게 넘겨줘야 하는 계주 경기와 같습니다.

이 논문은 이 경기를 완전히 바꿔놓는 Picard τ-leaping이라는 새로운 방법을 소개합니다. 이 방법이 어떻게 작동하는지 쉬운 비유를 통해 설명해 보겠습니다.

1. 기존 방식: 일렬로 늘어선 줄

기존의 방식은 영화관에 입장하기 위해 줄을 서 있는 사람들을 생각하면 됩니다. 한 사람이 티켓을 사고 안으로 들어간 후에야, 다음 사람이 티켓을 살 수 있습니다. 극장에 문이 100개 있더라도, 규칙상 "차례를 기다려야" 하기 때문에 한 번에 한 사람만이 계산대를 이용할 수 있습니다. 컴퓨터 용어로, 이것은 **순차적 샘플링(sequential sampling)**입니다. 정확하긴 하지만, 컴퓨터가 가진 전체 능력을 활용하지 못하기 때문에 매우 고통스러울 정도로 느립니다.

2. 새로운 아이디어: "시간 여행" 그룹

저자들은 줄이 하나씩 움직이기를 기다리는 대신, 시간의 덩어리를 하나의 블록으로 취급할 수 있다는 점을 깨달았습니다. 예를 들어, 다음 주의 날씨를 예측하고 싶다고 가정해 봅시다. 월요일, 화요일, 수요일을 하나씩 차례대로 계산하는 대신, 이렇게 말할 수 있습니다. "일주일 치 날씨를 한꺼번에 추측한 다음, 우리가 제대로 했는지 확인하고, 다시 한번 더 나은 결과로 다시 추측하자."

이것이 그들의 병렬 시간(Parallel-in-Time) 접근 방식의 핵심입니다. 그들은 시간의 블록(예: 재구성 과정의 10단계)을 가져와서, 컴퓨터의 수많은 코어를 사용하여 10단계를 동시에 해결하려고 시도합니다.

3. 비법: "피카르 반복(Picard Iteration)" (추측하고 확인하는 루프)

어떻게 10단계를 한꺼번에 해결하면서도 엉망이 되지 않을 수 있을까요? 저자들은 **피카르 반복(Picard iteration)**이라는 수학적 트릭을 사용합니다.

1라운드 (거친 추측): 컴퓨터는 시작점을 바탕으로 일주일 치 날씨(또는 전체 이미지 재구성)에 대한 대략적인 추측을 합니다.
2라운드 (수정): 컴퓨터는 게임의 "규칙"(수학 모델)을 살펴보고 첫 번째 추측이 어디서 틀렸는지 확인합니다. 컴퓨터 앞에 일주일 치 데이터가 통째로 놓여 있기 때문에, 모든 오류를 동시에 수정할 수 있습니다.
3라운드 (정교화): 이 과정을 반복합니다. 매번 추측은 진실에 더 가까워집니다.

컴퓨터가 "월요일부터 금요일까지"의 모든 수학적 계산을 정확히 동시에 수행할 수 있기 때문에, 일렬로 늘어선 줄 방식보다 훨씬 빠르게 작업을 마칠 수 있습니다.

4. 특별한 규칙: "퍼스트 히트(First-Hit)" 정지

여기에는 주의할 점이 있습니다. 이 특정 종류의 게임(흡수 확산, Absorbing Diffusion)에서는, 일단 "물음표"가 실제 글자나 픽셀로 변하면, 그것은 영원히 그 자리에 머뭅니다. 그것은 다시 변하지 않습니다.

만약 일주일 치를 한꺼번에 추측한다면, 이미 이전 단계에서 고정된 글자를 실수로 바꾸려고 할 수도 있습니다. 이를 해결하기 위해 저자들은 **"퍼스트 히팅 트렁케이션(First-Hitting Truncation)"**을 추가했습니다.

이것은 약간의 반전이 있는 "의자 뺏기" 게임과 같습니다: 일단 의자가 차지되면, 그 의자는 잠깁니다. 만약 당신의 "그룹 추측"이 이미 앉아 있는 사람을 움직이려 한다면, 시스템은 그 움직임을 무시하고 그들을 자리에 그대로 둡니다. 이는 컴퓨터가 속도를 높이려다 규칙을 어기는 일을 방지합니다.

5. 결과: 품질 손실 없는 속도

이 논문은 이 "덩어리로 추측하고 정교화하는" 방법을 사용함으로써 다음과 같은 성과를 냈다고 주장합니다:

속도: 단일 컴퓨터 칩(GPU)에서 기존 방식보다 품질을 똑같이 유지하면서도 이미지는 1.45배에서 1.86배 더 빠르게 생성할 수 있습니다.
효율성: 동일한 결과를 얻기 위해 약 **50% 적은 계산 단계(NFE)**가 필요합니다.
확장성: 이론적으로 문제가 커질수록(더 복잡한 이미지나 긴 텍스트), 이 방식은 기존 방식에 비해 상대적으로 더 빨라집니다.

요약

이 논문은 텍스트와 이미지를 생성하는 AI 모델을 실행하는 새로운 방법을 제시합니다. AI가 아주 작고 느린 단계를 차례대로 밟도록 강요하는 대신, AI가 시간을 가로질러 크고 병렬적인 도약을 할 수 있게 해줍니다. 그들은 큰 도약이 정확성을 유지하도록 "추측하고 확인하는" 루프를 사용하며, AI가 이미 고정한 부분을 망가뜨리지 않도록 "잠금" 규칙을 사용합니다. 그 결과, 디지털 콘텐츠를 만드는 더 빠르고 효율적인 방법을 만들어냈습니다.

기술 요약: 병렬 시간 샘플링을 통한 이산 확산 모델의 가속화

문제 정의

이산 확산 모델은 범주형 데이터를 위한 생성 모델링의 중심이 되었으며, 분자 설계, 단백질 공학, 텍스트 및 이미지 생성 분야에서 핵심적인 역할을 하고 있습니다. 지배적인 접근 방식인 **흡수 이산 확산(Absorbing Discrete Diffusion)**은 토큰을 특수한 마스크 상태로 점진적으로 교체하고 이 과정을 역전시키는 법을 학습합니다. 그러나 생성 과정은 본질적으로 순차적입니다. $\tau$ -leaping과 같은 기존의 순차적 샘플러는 현대적인 병렬 컴퓨팅 자원(예: GPU)을 충분히 활용하지 못해 효율적인 추론의 병목 현상을 일으킵니다. 핵심 질문은, 생성 품질을 희생하지 않으면서 병렬화가 이산 확산 모델의 샘플링을 근본적으로 가속화할 수 있는가 하는 것입니다.

방법론: Picard $\tau$ -Leaping

저자들은 흡수 이산 확산을 위한 최초의 병렬 시간 알고리즘인 Picard $\tau$ -leaping 방법을 제안합니다. 이 접근 방식은 연속 확산 모델에서 주로 사용되는 Picard 반복법을 이산 상태 공간에서의 $\tau$ -leaping 스토캐스틱 적분 형태로 확장한 것입니다.

핵심 알고리즘

시간 이산화: 역방향 샘플링 구간 $[\eta, T]$ 를 큰 블록들로 나눕니다. 각 블록은 다시 세밀한 시간 셀(time cells)로 하위 분할됩니다.
병렬 블록 처리: 전체 궤적을 단계별로 순차적으로 계산하는 대신, 알고리즘은 전체 블록에 대한 전이를 동시에 해결합니다.
Picard 반복:
- 초기화: 블록은 초기 추측값(블록 시작 시점의 상태)으로 시작합니다.
- 병렬 평가: 각 Picard 라운드에서, 블록 내의 모든 마이크로스텝(microsteps)은 이전 반복의 궤적을 사용하여 병렬로 평가됩니다.
- First-Hitting Truncation: 흡수 확산의 핵심 구성 요소입니다. 마스크된 토큰은 단 한 번만 언마스킹(unmasking)될 수 있으므로, 알고리즘은 블록 내에서 "first-hitting truncation"을 적용합니다. 이 연산은 각 좌표에 대해 가장 먼저 제안된 언마스킹 이벤트만을 적용함으로써, 병렬 접두사 연산(parallel prefix operations)을 가능하게 하는 동시에 흡수 구조를 보존합니다.
- 수렴: 프로세스는 수렴할 때까지 반복됩니다. 이론적 분석에 따르면, 오차는 Picard 반복 횟수에 따라 지수-팩토리얼(exponentially-factorially)적으로 수축합니다.

이론적 프레임워크

이 방법은 푸아송 무작위 측도(Poisson random measures)를 통해 이산 확산을 표현하는 것에 기반합니다. 저자들은 스코어 추정 오차, 유계된 스코어 추정치, 그리고 "정규화된 이벤트 스위칭(normalized event switching)" 조건(리프시츠 연속성과 유사함)에 관한 특정 가정을 도입하여 오차 전파를 제한합니다. 저자들은 Picard 반복이 순차적 $\tau$ -leaping 궤적으로 수렴함을 증명합니다.

주요 기여

1. 이론적 발전

병렬 시간 알고리즘: 흡수 이산 확산을 위한 Picard $\tau$ -leaping 방법(Algorithm 1)을 도입했습니다.
복잡도 감소: 함수 평가 횟수(NFE)에 대한 시간 복잡도가 주류 순차 방식의 $O(d \log S)$ 에서 $O(\log(d \log S) \cdot \log d)$ 로 개선됨을 증명했습니다.
수렴 증명: 이산 환경에서 Picard 반복이 갖는 **지수-팩토리얼 수렴(exponential-factorial convergence)**을 증명했습니다. 이는 일반적인 연속 확산의 지수적 수축과는 구별되는 특성입니다.
공간 복잡도: 블록 궤적을 저장하기 위해 $\tilde{O}(dS)$ 의 추가적인 공간 복합도를 부담하며, 이는 메모리를 병렬성을 위해 트레이드오프한 결과입니다.

2. 실증적 결과

저자들은 합성 데이터와 실제 작업에서 방법을 평가했습니다:

합성 데이터: 2D 분포(Chessboard 및 Circle)에서, 병렬 방법은 순차적 베이스라인과 유사한 KL 다이버전스를 유지하면서도 상당한 런타임 가속(특정 합성 스케일링 시나리오에서 최대 7~9배)을 달고 달성했습니다.
이미지 생성: ImageNet(256×256)에 대한 MaskGiT 기반 모델을 사용하여, 병렬 샘플러는 순차적 $\tau$ -leaping 베이스라인과 동등한 Fréchet Inception Distance (FID) 점수를 달성하면서도 50% 적은 NFE를 기록했습니다. 단일 RTX 4090 GPU에서 이는 1.86배의 벽 시계 시간(wall-clock) 가속으로 이어졌습니다.
텍스트 생성: OpenWebText(GPT-2 규모)에 대한 RADD 기반 모델을 사용하여, 병렬 방법은 고정된 NFE에서 더 나은 퍼플렉시티(perplexity)를 달성했으며, 순차적 솔버와 동일한 퍼플렉시티를 절반의 NFE만으로 달성했습니다. 단일 H100 GPU에서 이는 1.45배의 가속을 가져왔습니다.

의의 및 주장

본 논문은 이 연구가 효율적인 병렬 추론을 위한 이산 확산 모델의 잠재력을 확장한다고 주장합니다. 흡수 확산의 특정 구조적 특성(특히 언마스킹의 first-hitting 성질)을 활용함으로써, 저자들은 순차적 의존성을 깨뜨려 시간 지평(time horizon)을 따른 병렬 계산이 가능함을 입증했습니다.

저자들은 이를 좌표 수준의 병렬성(한 단계에서 여러 토큰을 언마스킹하는 것)에서 시간 수준의 병렬성(여러 타임 스텝을 동시에 해결하는 것)으로의 근본적인 전환으로 규정합니다. 저자들은 현재의 단일 GPU 가속이 메모리 트래픽에 의해 제한되지만, 폴리로그(polylogarithmic) 복잡도는 멀티 GPU 스케일링에서 훨씬 더 큰 잠재력을 시사한다고 언급합니다. 이 연구는 이산 생성 모델을 가속화하기 위한 이론적 토대를 제공하며, 분자 구조 생성 및 대규모 언어 모델링과 같이 빠른 샘플링을 요구하는 응용 분야에 시사점을 줍니다.

저자들은 현재의 분석이 총 변동(Total Variation, TV) 오차에 초점을 맞추고 있음을 인정하며, 향후 연구가 더 정교한 KL 수준의 분석과 더 높은 확장성을 확보하기 위한 더 메모리 효율적인 변형(예: 슬라이딩 윈도우 업데이트)을 목표로 해야 한다고 밝히며 겸허한 태도를 유지했습니다.