💻 computer science

Teacher-Feature Drifting: One-Step Diffusion Distillation with Pretrained Diffusion Representations

본 논문은 보조 네트워크의 필요성을 제거하기 위해 사전 훈련된 교사 모델의 자체 중간 은닉 상태를 특징 표현으로 활용하는 단순화된 1 단계 확산 증류 방법을 제안하며, 동시에 경쟁력 있는 FID 점수를 갖춘 고품질 다양성 이미지 생성을 달성하기 위해 경량 모드 커버리지 손실을 통합합니다.

원저자: Yuan Zhang, Chenyi Li, Guoqing Ma, Jiajun Zha, Yuanming Yang, Bo Wang, Wei Tang, Wenbo Li, Haoyang Huang, Nan Duan

게시일 2026-05-12

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Yuan Zhang, Chenyi Li, Guoqing Ma, Jiajun Zha, Yuanming Yang, Bo Wang, Wei Tang, Wenbo Li, Haoyang Huang, Nan Duan

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

마스터 셰프 (Teacher) 가 있다고 상상해 보세요. 이 셰프는 가장 맛있고 고품질의 요리를 만드는 것으로 유명합니다. 하지만 이 셰프는 놀라울 정도로 느립니다. 완벽한 한 끼 식사를 만들기 위해 서빙하기 전까지 재료를 50 번 또는 100 번씩 맛보고, 조정하고, 다듬어야 합니다. 여러분은 이만큼 훌륭한 요리를 만들 수 있지만, 단 한 번의 단계로 만들어내는 학생 셰프 (Student) 를 원합니다.

이 논문은 그 학생 셰프를 가르치는 새로운 방법을 소개합니다. 이를 **Teacher-Feature Drifting(TFD)**이라고 부릅니다. 간단한 비유를 통해 작동 원리를 설명해 보겠습니다.

1. 문제: 단계가 너무 많습니다

일반적으로 학생이 마스터처럼 요리하도록 가르치려면, 학생이 마스터가 요리하는 과정을 단계별로 지켜보게 하거나, 별도의 '맛 평가' 로봇을 고용해 학생의 요리를 채점하게 할 수 있습니다. 이러한 방법들은 복잡하고, 추가 장비를 필요로 하거나, 훈련에 시간이 오래 걸립니다.

2. 해결책: 마스터 자신의 '내부 나침반' 활용

저자들은 마스터 셰프 (Pretrained Diffusion Model) 가 이미 뇌 안에 내장된 '맛 감각'을 가지고 있음을 깨달았습니다. 요리하는 동안에도 마스터의 뇌는 다지기, 섞기, 끓이기 등 음식의 다양한 단계를 처리합니다.

별도의 맛 평가 로봇을 고용하는 대신, TFD 는 마스터 자신의 뇌 상태를 측정 도구로 사용합니다.

비법: 학생이 요리를 시도할 때, 시스템은 마스터가 같은 요리를 만들고 있었다면 그 정확히 그 순간 마스터의 뇌가 무엇을 생각했는지 살펴봅니다.
'Drifting(이동)'의 의미: 학생의 요리를 배라고 상상해 보세요. 마스터의 뇌는 배를 '완벽한 요리' 쪽으로 부드럽게 밀어주고 '나쁜 요리' 쪽에서는 밀어내는 해류를 만듭니다. 학생은 단지 해류가 지시하는 방향으로 배를 조종하기만 하면 됩니다.

3. 비밀 재료: 약간의 '노이즈'

이 논문은 놀라운 사실을 발견했습니다. 마스터 셰프에게 완벽하게 깨끗하고 완성된 요리를 평가하게 하면, 평가가 너무 날카롭고 까다로워집니다. 마치 현미경으로 그림을 보며 평가하는 것과 같습니다. 미세한 먼지 입자에 주의가 산만해질 수 있습니다.

대신 저자들은 약간 흐릿하거나 '노이즈'가 섞인 요리 (아직 초점이 완전히 맞지 않은 사진과 같은) 를 마스터에게 보여주는 것이 가장 효과적임을 발견했습니다.

이유: 이 '흐림'은 마스터의 판단을 부드럽게 만듭니다. 학생이 사소한 중요하지 않은 세부 사항에 집착하는 것을 막고, 큰 그림 (형태, 색상, 전체적인 분위기) 에 집중하도록 돕습니다. 이로 인해 학습 과정이 훨씬 매끄러워지고 최종 결과물도 더 좋아집니다.

4. '모방' 문제 피하기

AI 교육에서 흔한 문제는 **Mode Collapse(모드 붕괴)**입니다. 이는 한 가지 완벽한 피자를 만드는 법을 배운 학생 셰프가, 고객이 샐러드를 주문했음에도 불구하고 모든 주문에 대해 정확히 같은 피자를 만들기로 결정하는 것과 같습니다. 학생은 다양성을 탐색하는 것을 멈춥니다.

이를 해결하기 위해 저자들은 **"Coverage Loss(Anchor-Margin Loss)"**를 추가했습니다.

비유: 마스터 셰프가 만들 수 있는 모든 맛있는 음식에 대한 지도를 가지고 있다고 상상해 보세요. 학생에게 이렇게 말합니다. "지도의 한 곳에만 서 있지 마세요. 지도의 서로 다른 지역을 모두 방문하도록 하세요."
시스템은 확인합니다. "학생이 마스터 지도의 특정 부분을 커버하는 요리를 만들어 보았는가?" 만약 학생이 어떤 지역을 무시한다면, 시스템은 그 지역으로 가도록 밀어줍니다. 이를 통해 학생이 한 가지 유형이 아닌 다양한 요리를 만들 수 있도록 학습하게 됩니다.

결과

이 논문은 두 가지 큰 도전 과제에서 이 방법을 테스트했습니다.

ImageNet: 1,000 가지 다른 사물의 이미지를 생성합니다. 새로운 방법은 단 한 단계로 고품질 이미지를 생성했는데, 이는 느린 50 단계 마스터 셰프의 결과와 거의 비슷했으며, 다른 빠른 방법들보다 훨씬 좋았습니다.
SDXL(텍스트-to-이미지): "모자를 쓴 고양이"와 같은 단어를 이미지로 변환합니다. 역시 새로운 방법은 고품질과 좋은 다양성을 유지하면서 단 한 단계로 수행했습니다.

요약

간단히 말해, 이 논문은 다음과 같이 말합니다. "학생을 가르치기 위해 새로운 교사를 고용하지 마세요. 마스터 자신의 내부 생각을 가이드로 활용하고, 수업을 더 쉽게 이해할 수 있도록 약간의 '흐림'을 추가하며, 학생이 한 가지 요리가 아닌 전체 메뉴를 탐색하도록 하세요."

이로써 훈련은 빠르고, 간단하며, 효과적이 되어 복잡한 추가 도구가 필요 없게 됩니다.

기술 요약: Teacher-Feature Drifting (TFD)

문제 정의

사전 훈련된 확산 (diffusion) 및 흐름 매칭 (flow-matching) 모델은 고충실도 이미지 생성을 달성하지만, 종종 단일 이미지를 생성하는 데 수십에서 수백 번의 순전파 (forward passes) 가 필요하여 추론 비용이 매우 높다는 단점이 있습니다. 이러한 다단계 교사 모델을 효율적인 1 단계 또는 소수 단계 생성기로 압축하기 위한 증류 (distillation) 방법이 존재하지만, 현재 접근 방식들은 상당한 복잡성에 직면해 있습니다. 분포 매칭 증류 (Distribution Matching Distillation, DMD) 와 그 변형들과 같은 방법들은 보조 네트워크 (예: 가짜 점수 네트워크, 적대적 판별기) 나 다단계 훈련 일정에 의존합니다. Consistency Models 와 같은 다른 접근 방식들은 불안정한 훈련을 겪거나 복잡한 야코비 - 벡터 곱 (Jacobian-vector products) 을 필요로 할 수 있습니다. 또한, 최근의 "Drifting Model" 목적 함수들은 1 단계 생성을 위한 직접적인 인력 - 반발력 (attraction-repulsion) 손실률을 제공하지만, 일반적으로 표현 공간을 정의하기 위해 외부 사전 훈련된 특징 추출기 (예: 자기지도 학습 모델 또는 latent-MAE) 에 의존합니다. 이는 추가적인 훈련 오버헤드와 아키텍처 복잡성을 초래합니다.

해결되는 핵심 문제는 보조 네트워크, 복잡한 최적화 파이프라인, 또는 외부 특징 추출기에 의존하지 않고 단일 목적 함수를 사용하여 효율적이고 고품질의 1 단계 확산 증류를 달성하는 방법입니다.

방법론: Teacher-Feature Drifting (TFD)

저자들은 사전 훈련된 교사 모델 자체의 내부 표현을 활용하여 1 단계 증류를 단순화하는 Teacher-Feature Drifting (TFD) 프레임워크를 제안합니다.

1. Teacher-Induced 표현 공간

외부 특징 인코더가 필요한 원래 Drifting Model 과 달리, TFD 는 드리프트 손실 (drifting loss) 을 계산하기 위한 특징 공간으로 고정된 사전 훈련된 교사 모델의 **중간 은닉 상태 (intermediate hidden states)**를 사용합니다.

메커니즘: 생성된 샘플 $x$ 와 실제 샘플 $x^+$ 에 대해, 교사 네트워크 $T$ 의 특정 층 $\ell$ 에서 특징을 추출합니다: $\phi^\ell_T(x, c) = h^\ell_T(x, t_f, c)$ .
Drifting Loss: 드리프트 필드는 이 교사 유도 공간에서 직접 계산됩니다. 생성기 $f_\theta$ 는 교사 특징 공간의 인력 - 반발력 필드에 의해 정의된 "고정된 운송 목표 (frozen transported target)"와 출력 간의 거리를 최소화하도록 최적화됩니다. 이는 증류 목적 함수를 증류되는 모델과 정렬시켜 별도의 표현 네트워크가 필요 없게 합니다.

2. 노이즈가 추가된 특징 추출

중요한 경험적 발견은 적당히 노이즈가 추가된 입력에서 특징을 추출하는 것이 깨끗한 샘플을 사용하는 것보다 더 좋은 결과를 낸다는 것입니다.

샘플을 특징 추출을 위해 교사 모델에 입력하기 전에, 교사의 순전파 노이즈 과정으로 교란 (perturb) 시킵니다: $x_{tf} = x + \sigma_{tf} \xi$ .
저자들은 이 교란이 교사 특징 기하학을 매끄럽게 하여 고주파 불일치를 억제하고 분포 수준 매칭에 더 적합한 표현 공간을 만든다고 가정합니다. 실험에 따르면 노이즈 수준 $\sigma_{tf} = 0.1$ 이 최적입니다.

3. Anchor-Margin 정규화

모드 붕괴 (mode collapse) 문제 (생성기가 교사가 지원하는 다양한 영역을 포괄하지 못하는 현상) 를 해결하기 위해, 저자들은 경량 anchor-margin loss를 도입합니다.

Anchors: 실제 또는 교사 생성 샘플에서 일련의 앵커 특징 $\{a_i\}$ 를 추출합니다.
메커니즘: 이 손실은 생성된 샘플이 이러한 앵커 주변의 영역을 차지하도록 장려합니다. 이는 앵커 분포 자체의 국부적 밀도에 기반하여 각 앵커에 필요한 "지지 (support)"를 적응적으로 조정합니다.
목표: 이는 지역적 인력 - 반발력 드리프트 목적 함수에 보완적인 커버리지 신호를 제공하여, 새로운 네트워크를 도입하지 않고도 학생 생성기가 다양한 교사 지원 영역을 포괄하도록 보장합니다.

최종 목적 함수는 Teacher-Feature Drifting 손실 ( $L_{TFD}$ ) 과 anchor-margin 손실 ( $L_{anchor}$ ) 을 결합합니다:
$L = L_{TFD} + \lambda_{anchor} L_{anchor}$

주요 기여

단순화된 프레임워크: TFD 는 사전 훈련된 교사의 은닉 표현 공간에서 드리프트 목적 함수를 직접 계산하는 1 단계 확산 증류 프레임워크를 제안하여 추가 표현 네트워크의 필요성을 제거합니다.
노이즈가 추가된 특징 추출: 이 연구는 적당히 노이즈가 추가된 입력에서 교사 특징을 추출하면 증류를 위한 특징 기하학이 개선되어 더 좋은 FID 점수로 이어진다는 것을 보여줍니다.
모드 붕괴 완화: 샘플 다양성을 장려하고 교사 지원 특징 영역의 커버리지를 확보하기 위해 경량 anchor-margin 손실이 도입되어 증류에서 흔히 발생하는 실패 모드를 해결합니다.
성능: 이 방법은 단일 순전파로 경쟁력 있는 이미지 품질과 다양성을 달성하며, 최첨단 기준 모델에 비해 전체 증류 파이프라인을 실질적으로 단순화합니다.

실험 결과

저자들은 ImageNet-64×64(클래스 조건부) 와 SDXL(텍스트 - 이미지) 생성에서 TFD 를 평가했습니다.

ImageNet-64×64:
- TFD 는 20 만 번의 훈련 반복 후 단일 순전파로 FID 1.58을 달성했습니다.
- 이 결과는 DMD2(FID 1.51) 와 경쟁력 있으며, DMD(2.62), EMD(2.20), CTM(1.92) 와 같은 다른 1 단계 기준 모델보다 훨씬 뛰어납니다.
- 수렴: TFD 는 1.0k 업데이트에서 FID $\le$ 10 에 도달한 반면, DMD2 기준 모델은 5.5k 업데이트가 필요했습니다 (5.5 배 빠름). FID $\le$ 3 에 도달하는 데 TFD 는 12.5k 업데이트가 소요된 반면 DMD2 는 24.0k 업데이트가 필요했습니다.
SDXL 텍스트 - 이미지:
- COCO 2014 데이터셋에서 TFD 는 단일 순전파로 FID 18.39를 달성하여 다른 증류된 SDXL 변형들 (예: DMD2 의 19.01, SDXL-Turbo 의 24.57) 보다 우수했습니다.
- 이는 강력한 텍스트 - 이미지 정합성을 나타내는 경쟁력 있는 CLIP 점수 0.332 를 유지했습니다.
- 정성적 결과는 TFD 가 4 단계 증류 기준 모델과 50 단계 교사 모델과 비교할 수 있는 일관된 구조, 시각적 세부 사항, 그리고 프롬프트 정합성을 유지하면서도 단 한 단계만 사용함을 보여주었습니다.

중요성과 주장

이 논문은 TFD 가 기존 증류 프레임워크에 비해 더 단순하고 효율적인 대안을 제공한다고 주장합니다. 교사 모델 자체를 특징 추출기로 활용함으로써 TFD 는 보조 네트워크 훈련의 복잡성이나 다단계 파이프라인 설계를 피합니다. 노이즈가 추가된 특징 추출과 anchor-margin 정규화의 도입은 각각 표현 기하학과 모드 커버리지의 특정 문제를 해결합니다.

저자들은 TFD 를 경쟁력 있는 이미지 품질과 다양성을 갖춘 효율적인 1 단계 생성을 달성하는 방법으로 제시하며, 교사의 고유 표현 공간에서의 단일 목적 함수가 고충실도 증류를 위해 충분함을 입증합니다. 이 연구는 교사의 내부 상태를 올바르게 활용한다면 외부 특징 인코더 없이도 "드리프트" 개념을 효과적으로 실현할 수 있음을 시사합니다.

한계점: 저자들은 현재 작업이 U-Net 기반 교사 모델에 초점을 맞추고 있으며, 아직 Diffusion Transformer (DiT) 아키텍처를 탐구하지 않았다고 지적합니다. 또한, 특징 층과 노이즈 수준의 선택은 현재 이론적 지침이 부족하며, 이 방법은 현재 실험을 넘어 더 큰 규모의 데이터셋이나 더 높은 해상도 설정에서 테스트되지 않았습니다.