🤖 AI

Infinite Mask Diffusion for Few-Step Distillation

본 논문은 표준 마스킹 확산 모델의 이론적 인수분해 오차 한계를 극복하기 위해 확률적 무한 상태 마스를 활용하는 무한 마스크 확산 모델 (IMDM) 을 제안하며, 이를 통해 효율적인 소수 단계 증류와 소수 단계 텍스트 생성 작업에서의 우수한 성능을 가능하게 한다.

원저자: Jaehoon Yoo, Wonjung Kim, Chanhyuk Lee, Seunghoon Hong

게시일 2026-05-12

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Jaehoon Yoo, Wonjung Kim, Chanhyuk Lee, Seunghoon Hong

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

"Few-Step Distillation 을 위한 Infinite Mask Diffusion"논문에 대한 설명을 쉬운 언어와 창의적인 비유로 제시합니다.

큰 그림: "빈칸 채우기" 게임

일부 단어가 검은 상자 (마스크) 뒤에 숨겨진 문장을 완성해야 하는 게임을 상상해 보세요.

옛 방식 (자기회귀 모델): 첫 번째 숨겨진 단어를 추측한 다음, 두 번째, 세 번째를 하나씩 차례로 추측해야 합니다. 한 번에 한 페이지만 넘길 수 있는 책을 읽는 것과 같습니다. 정확하지만 느립니다.
새 방식 (마스크 확산 모델 - MDMs): 모든 숨겨진 단어를 동시에 추측할 수 있습니다. 전체 페이지를 보고 모든 빈칸을 동시에 채우는 것과 같습니다. 이는 훨씬 빠르며, 모델이 한 번에 전체 문장의 맥락을 이해할 수 있게 합니다.

문제: "일률적"인 마스크

이 논문은 현재의 "새 방식"(MDMs) 에 있는 치명적인 결함을 지적합니다.

비유:
단어를 가리는 검은 상자가 모두 정확히 같은 재질, 즉 단일하고 고체이며 불투명한 검은 플라스틱 시트로 만들어져 있다고 상상해 보세요.
모델이 상자 아래에 무엇이 있는지 추측하려 할 때, 그 "플라스틱 시트"가 단어들 간의 관계를 어떻게 연결하는지에 대한 단서를 전혀 제공하지 않기 때문에, 모델은 모든 단어를 독립적으로 추측해야 합니다. 마치 100 개의 분리되고 동일한 검은 스크린을 보며 영화의 줄거리를 추측하려는 것과 같습니다.

모델이 연결 관계를 보지 못한 채 모두 한 번에 추측하기 때문에 **분해 오차 (Factorization Error)**라는 특정 유형의 실수를 범합니다. 이는 조각들이 잘못된 순서로 붙어 있는 퍼즐을 풀려는 것과 같습니다. 이를 고치기 위해 모델은 보통 여러 번 추측하고, 확인하고, 다시 추측하는 과정 (반복 단계) 을 거쳐야 하므로, 빠르다는 목적 자체가 무너집니다.

저자들은 이론적 한계를 발견했습니다. 모델이 아무리 똑똑해지더라도, 그 단일하고 고체인 검은 플라스틱 시트를 사용한다면 단 한두 단계로 단어를 완벽하게 추측할 수 없습니다. 실수가 얼마나 나쁠 수 있는지에 대한 "바닥"이 존재하며, 이는 빠른 생성에는 너무 높은 장벽입니다.

해결책: "무지개 무한" 마스크

저자들은 **IMDM(무한 마스크 확산 모델)**이라는 새로운 모델을 제안합니다.

비유:
단일한 고체 검은 플라스틱 시트 대신, 마스크가 무한하고 고유하며 투명한 유색 유리로 만들어졌다고 상상해 보세요.

단어가 숨겨질 때마다 고유하고 무작위인 "색상"이나 "질감"(확률적 잠재 마스크) 을 부여받습니다.
이러한 마스크들이 서로 다르게 보이지만, 모델은 여전히 이를 실제 단어와 구별할 수 있습니다.
숨겨진 자리마다 고유한 "지문"이 있기 때문에, 모델은 그 무작위성을 활용하여 숨겨진 단어들이 서로 어떻게 연결되는지 파악할 수 있습니다.

이는 모델에게 숨겨진 단어들 사이의 보이지 않는 연결고리를 볼 수 있게 해주는 특별한 안경을 주는 것과 같습니다. 이러한 "무한한 다양성"의 마스크를 사용함으로써 모델은 기존 모델들을 가두었던 실수의 "바닥"을 우회할 수 있습니다.

마법 같은 트릭: 증류 (Distillation)

논문은 증류에 대해서도 언급합니다. 이는 교사 - 학생 관계를 생각하면 됩니다.

교사: 정답을 얻는 데 100 단계가 걸리는 느리지만 완벽한 모델.
학생: 단 2 또는 4 단계로 수행하도록 배워야 하는 빠른 모델.

보통 학생은 "검은 플라스틱 시트" 문제 (오차 바닥) 로 인해 복잡한 연결 관계를 빠르게 배우는 데 실패합니다. 하지만 새로운 IMDM 모델이 "무지개 무한 마스크"를 사용하기 때문에, 학생은 실제로 교사의 비밀을 배울 수 있습니다. 몇 단계 만에 교사의 복잡하고 다단계적인 사고 방식을 모방할 수 있게 됩니다.

논문이 발견한 것

이론: 그들은 수학적으로 기존 모델 (MDMs) 이 빠르기를 시도할 때 최소한의 오차 수준에 갇혀 있음을 증명했습니다. 새로운 모델 (IMDM) 은 이 한계를 제거합니다.
합성 테스트: 두 단어가 반드시 동일해야 하는 ("00"또는"11"과 같은) 간단한 퍼즐을 만들었습니다. 기존 모델은 무작위로 추측 (50/50) 하여 실패했습니다. 반면 새로운 모델은 한 단계 만에 거의 100% 의 정확도로 맞췄습니다.
실제 세계 테스트: 그들은 LM1B 와 OpenWebText 와 같은 거대한 텍스트 데이터셋에서 이를 테스트했습니다.
- 매우 적은 단계(2~8 단계)로 텍스트를 생성하도록 요청했을 때, 새로운 IMDM 모델은 기존 방법보다 훨씬 더 고품질의 텍스트를 생성했습니다.
- 기존 모델은 서두르면 의미 없는 말이나 반복적인 텍스트를 생성했습니다. 반면 새로운 모델은 일관성 있고 다양하며 문법적으로 올바른 문장을 생성했습니다.

요약

논문의 결론은 다음과 같습니다: "우리는 현재 텍스트 생성의 빠른 방식이 단일하고 지루한 마스크를 사용함으로써 숨겨진 속도 저하를 겪고 있음을 발견했습니다. 우리는 모델에 고유하고 무작위적인 무한한 다양성의 마스크를 제공함으로써 이를 수정했습니다. 이를 통해 모델은 단어들이 서로 어떻게 연결되는지를 즉시 이해할 수 있게 되어, 수십 단계 대신 단 몇 단계 만에 고품질 텍스트를 생성할 수 있게 되었습니다."

기술 요약: 소수 단계 증류를 위한 무한 마스크 확산

문제 정의

마스크 확산 모델 (MDM) 은 병렬 디코딩과 양방향 문맥 처리와 같은 장점을 제공하며, 언어 모델링을 위한 자기회귀 (AR) 모델의 유망한 대안으로 부상했습니다. 그러나 MDM 은 일반적으로 고품질 텍스트를 생성하기 위해 많은 샘플링 반복이 필요합니다. 이러한 필요성은 여러 토큰이 동시에 업데이트될 때 발생하는 분해 오차에서 비롯됩니다.

저자들은 표준 MDM 의 근본적인 이론적 한계를 규명했습니다: 결정론적 단일 상태 마스크 토큰(예: [MASK])의 사용은 분해 오차에 대한 회복 불가능한 하한을 생성합니다. 마스크 토큰이 결정론적이며 모든 마스크된 위치에 대해 동일하기 때문에, 모델은 마스크된 상태의 서로 다른 확률적 실현을 구별할 수 없습니다. 결과적으로, 소수 단계 생성에 필요한 여러 상관 토큰이 동시에 디코딩될 때, 모델은 그들의 결합 의존성을 포착하지 못하여 어떤 증류 방법을 사용하더라도 성능의 한계가 지속됩니다. 심지어 고급 소수 단계 증류 기술을 사용하더라도 표준 MDM 은 극히 적은 단계 (예: 2~4 단계) 영역에서 의미적으로 일관된 데이터를 생성하는 데 어려움을 겪습니다.

방법론: 무한 마스크 확산 모델 (IMDM)

이 구조적 병목 현상을 해결하기 위해 저자들은 **무한 마스크 확산 모델 (IMDM)**을 제안합니다. 핵심 혁신은 결정론적 단일 상태 마스크를 확률적 무한 상태 마스크로 대체하는 것입니다.

주요 기술 구성 요소:

확률적 무한 상태 마스크:
- 단일 마스크 토큰 $m$ 대신, IMDM 은 $|\mathcal{M}| \to \infty$ 인 잠재 마스크 공간 $\mathcal{M}$ 을 정의합니다.
- 순방향 과정은 마스크된 토큰을 이 무한 집합에서 균일하게 무작위로 선택된 잠재 마스크 상태로 전이시킵니다.
- 결정적으로, 모델은 MDM 의 구별 가능성 속성을 유지합니다: 데이터 토큰 집합 $\mathcal{V}$ 와 마스크 상태 집합 $\mathcal{M}$ 은 엄격하게 분리되어 있습니다 ( $\mathcal{V} \cap \mathcal{M} = \emptyset$ ). 이는 모델이 여전히 관측된 위치와 예측해야 할 위치를 명시적으로 식별할 수 있음을 보장합니다.
연속 노이즈를 통한 구현:
- 무한 집합을 명시적으로 인스턴스화하는 것은 계산적으로 불가능하므로, 저자들은 마스크 임베딩에 연속 노이즈 $\epsilon$ 을 주입하여 IMDM 을 구현합니다.
- 마스크 임베딩은 $E_\theta(z) = E_\psi(m) + \text{MLP}_\theta(\epsilon)$ 로 표현되며, 여기서 $\epsilon$ 은 균일 분포에서 샘플링됩니다.
- 이 설계는 모델이 무한한 종류의 마스크를 시뮬레이션하면서도 사전 훈련된 MDM 가중치와 호환되도록 합니다. MLP 는 0 으로 초기화되어 훈련 시작 시 모델이 사전 훈련된 MDM 과 동일하게 작동하도록 보장합니다.
분해 오차의 이론적 완화:
- 저자들은 표준 MDM 이 조건부 총 상관관계 (분해 오차) 에 대해 0 이 아닌 하한을 갖는 반면, IMDM 은 이 오차를 0 으로 줄일 수 있는 최적 모델을 허용함을 증명합니다.
- 이는 "분할 - 매핑" 메커니즘을 통해 달성됩니다: 마스크 상태의 무한한 확률성은 모델이 노이즈 공간을 특정 토큰 시퀀스에 해당하는 부분 집합으로 분할할 수 있게 합니다. 그런 다음 모델은 이러한 노이즈 실현을 일관된 토큰 시퀀스로 결정론적으로 매핑하여, 결정론적 마스크가 포착하지 못하는 결합 의존성을 효과적으로 포착할 수 있습니다.
증류 호환성:
- IMDM 은 기존 소수 단계 증류 프레임워크 (예: SDTT, ReDi, Di4C) 와 원활하게 통합하도록 설계되었습니다.
- 훈련 목적 함수 (NELBO) 는 표준 MDM 과 동일한 형태를 유지하여 사전 훈련된 가중치의 직접적인 이전과 확립된 증류 전략의 적용을 용이하게 합니다.

주요 기여

분해 오차에 대한 이론적 분석: 이 논문은 MDM 에 대한 분해 오차의 이론적 하한을 확립하여, 그것이 단일 상태 마스크의 결정론적 성격에서 비롯되며 증류 방법과 무관하게 지속됨을 보여줍니다.
IMDM 제안: MDM 의 간단한 프레임워크와 조건부 생성 능력을 유지하면서 이 이론적 하한을 완화하는 확률적 무한 상태 마스크를 도입했습니다.
실증적 검증:
- 합성 작업: 완벽한 토큰 상관관계가 필요한 합성 데이터셋 (이진 시퀀스 $\{00, 11\}$ ) 에서 표준 MDM 은 단일 단계에서 유효한 샘플을 생성하지 못했습니다 (유효성 $\approx$ 50%), 반면 IMDM 은 근본적인 의존성 구조를 학습하여 거의 완벽한 유효성 (97.7%) 을 달성했습니다.
- 언어 모델링 벤치마크: LM1B 와 OpenWebText 에서 IMDM 은 증류 방법 (특히 SDTT+ReDi) 과 결합되어 소수 단계 영역 ( $\le 8$ 단계) 에서 기존 소수 단계 베이스라인 (MDLM, SDTT, Di4C) 을 크게 능가했습니다.
- 지표: IMDM 은 특히 낮은 단계 수 (2~4 단계) 에서 조건부 생성 작업에서 베이스라인에 비해 더 낮은 생성 퍼플렉시티 (Gen. PPL) 와 더 높은 MAUVE 점수를 달성했습니다.

결과

소수 단계 성능: IMDM 은 소수 단계 영역에서 베이스라인에 비해 상당한 성능 차이를 보여주었습니다. 예를 들어, ReDi 를 사용한 2 단계 생성에서 LM1B 에 대해 IMDM 은 증류된 MDM 베이스라인에 비해 Gen. PPL 을 거의 50% 감소시켰습니다.
확장성: 성능 향상은 다양한 모델 규모 (작은 모델에서 8 억 6 천만 파라미터 모델까지) 와 데이터셋 (LM1B 및 OpenWebText) 에서 일관되었습니다.
절대 연구: 저자들은 주입된 노이즈의 차원이 중요하다는 것을 발견했습니다; 768 차원은 LM1B 에 충분했지만, 마스크 임베딩에서 충분한 구별성을 보장하기 위해 더 복잡한 OpenWebText 데이터셋에는 2048 차원이 필요했습니다. 이 방법은 노이즈 분포 (균일 대 가우시안) 와 스케일 선택에 대해 견고했습니다.

중요성과 주장

이 논문은 IMDM 이 효율적인 소수 단계 언어 생성 개발을 이전에 방해했던 MDM 프레임워크에 내재된 구조적 한계를 해결한다고 주장합니다. 마스크 토큰에 확률성을 도입함으로써 IMDM 은 모델이 회복 불가능한 분해 오차 하한에서 해방되어, 극히 큰 단계 크기 (소수 단계) 로도 토큰 의존성을 효과적으로 포착할 수 있게 합니다.

저자들은 IMDM 을 MDM 프레임워크의 대체제가 아니라, 저지연 추론을 위한 그 잠재력을 완전히 해제하기 위한 필수적인 진화로 위치시킵니다. 이 연구는 적절한 아키텍처 수정 (무한 확률적 마스크) 과 증류 전략을 통해 확산 기반 언어 모델이 소수 단계 생성 시나리오에서 자기회귀 모델과 경쟁하거나 더 우수한 성능을 달성할 수 있음을 보여줍니다. 이는 MDM 이 이전에 어려움을 겪었던 영역입니다. 이 방법은 MDM 의 단순성을 유지하여 사전 훈련된 가중치의 재사용과 기존 훈련 파이프라인과의 호환성을 가능하게 합니다.