💻 computer science

DiverseDiT++: Quantifying, Analyzing, and Promoting Representation Diversity in Diffusion Transformers

이 논문은 디퓨전 트랜스포머(Diffusion Transformer)의 블록 단위 표현 다양성을 정량화하기 위해 가중 다양성 점수(Weighted Diversity Score, WDS)를 도입하여 이것이 합성 품질과 강한 상관관계가 있음을 밝히고, 긴 잔차 연결(long residual connections)과 다양성 손실(diversity loss)을 통해 다양한 표현 학습을 명시적으로 촉진함으로써 성능을 향상시키는 프레임워크인 DiverseDiT++를 제안한다.

원저자: Binglei Li, Mengping Yang, Zhiyu Tan, Xiaomeng Yang, Zhizhong Huang, Junping Zhang, Hao Li

게시일 2026-08-05

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Binglei Li, Mengping Yang, Zhiyu Tan, Xiaomeng Yang, Zhizhong Huang, Junping Zhang, Hao Li

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

당신이 로봇에게 걸작을 그리는 법을 가르치려 한다고 상상해 보세요. 당신은 단순히 로봇이 그림을 똑같이 베끼는 것을 원하는 것이 아니라, 로봇이 세상을 아주 잘 이해하여 무에서 유로 새롭고 아름다운 장면을 창조해내기를 원합니다. 이것이 바로 기계가 예술, 음악, 심지어 분자까지도 스스로 만들어내는 법을 배우는 컴퓨터 과학의 한 분야인 '생성형 AI(generative AI)'의 목표입니다. 이를 위해 현대의 AI는 '확산 모델(diffusion model)'이라는 영리한 기술을 사용합니다. 이것은 마치 노이즈와 정전기 가득한 찰흙 덩어리에서 시작하여, 완벽한 조각상이 나타날 때까지 노이즈를 조금씩 깎아 나가는 조각가와 같습니다. 하지만 이 조각상이 실제처럼 보이고 상세하게 표현되려면, AI에게 생각을 정리할 수 있는 두뇌가 필요합니다. 최근 과학자들은 기존의 단순한 두뇌를 인간의 언어를 이해하는 데 도움을 주는 강력한 구조인 '트랜스포머(Transformers)'로 교체했습니다. 이 새로운 '확산 트랜스포머(Diffusion Transformers)'는 그림을 그리는 데 매우 뛰어나지만, 과학자들은 여로 여전히 이들의 두뇌 내부에서 정확히 어떻게 작동하는지를 알아내기 위해 노력하고 있습니다. 큰 질문은 이것입니다. 무엇이 이 AI 예술가들을 이토록 재능 있게 만드는 걸까요? 단지 더 큰 두뇌를 갖는 것의 문제일까요, 아니면 그들이 생각하는 방식에 숨겨진 비결이 있는 것일까요?

**DiverseDiT++**라는 제목의 이 논문은 이 모델들이 내부적인 생각을 어떻게 조직하는지를 살펴봄으로써 그 미스터리를 파헤칩니다. 연구진은 그 비결이 단순히 더 많은 데이터나 더 큰 모델을 갖는 것이 아니라, 바로 **다양성(diversity)**에 있다는 것을 발견했습니다. 벽화를 그리는 예술가 팀을 상상해 보세요. 만약 모든 예술가가 똑같은 것을 똑같은 방식으로 그리라는 지시를 받는다면, 최종 결과물은 지루하고 평면적일 것입니다. 하지만 각 예술가가 장면의 서로 다른 부분에 집중하도록 권장받는다면—한 명은 하늘에, 한 명은 나무에, 한 명은 그림자에 집중한다면—그 결과는 풍부하고 복잡한 걸작이 될 것입니다. 저자들은 확산 트랜스포머도 이와 똑같이 작동한다는 것을 발견했습니다. 학습함에 따라, 두뇌의 서로 다른 층(layer)들은 자연스럽게 전문화되기 시작하며 고유한 역할을 맡게 됩니다. 그러나 논문은 이러한 자연스러운 과정이 항상 충분히 강력하지는 않다고 제안합니다.

이를 증명하기 위해 연구팀은 **가중 다양성 점수(Weighted Diversity Score, WDS)**라는 새로운 측정 도구를 만들었습니다. 이것은 AI의 여러 층 사이의 생각이 얼마나 다른지를 체크하는 '창의성 측정기'와 같습니다. 그들은 강력한 연결 고리를 발견했습니다. 이 창의성 측정치가 높을 때마다 AI는 더 낫고 사실적인 이미지를 생성했습니다. 실제로 그들은 97가지의 서로 다른 실험을 통해 높은 다양성 점수와 더 나은 이미지 품질 사이에 매우 밀접한 관계(상관관계 -0.869)가 있음을 발견했습니다. 이는 만약 당신이 AI의 두뇌 층들이 서로 더 다르게 행동하도록 강제할 수 있다면, AI가 자신의 일을 더 잘하게 된다는 것을 시사합니다.

이 발견을 바탕으로 연구진은 **DiverseDiT++**라는 새로운 프레임워크를 구축했습니다. AI에게 무엇을 할지 가르치기 위해 외부의 조력자나 값비싼 사전 학습된 모델에 의존하는 대신, 그들은 AI의 자체 구조를 수정했습니다. 그들은 '긴 잔차 연결(long residual connections)'을 추가했는데, 이는 뇌의 초기 층이 나중 층과 직접 대화할 수 있게 하여 정보가 정체되거나 반복되지 않도록 하는 초고속 고속도로와 같습니다. 또한 그들은 특별한 '다양성 손실(diversity loss)' 규칙을 추가했습니다. 이것은 층들에게 "이봐, 너희 너무 똑같은 것만 하고 있어! 다른 걸 시도해 봐!"라고 말하는 엄격한 미술 선생님처럼 작동합니다. 이는 각 층이 고유한 특징을 학습하도록 강제합니다.

결과는 인상적이었습니다. DiverseDiT++를 다양한 크기의 AI 모델에 적용했을 때, 이미지는 더 선명해졌고 모델은 더 빠르게 학습되었습니다. 연구진은 이를 256 × 256 및 512 × 512 해상도의 ImageNet과 같은 표준 이미지 데이터셋에서 테스트했으며, 새로운 방식이 기존 방식들을 일관되게 앞질렀습니다. AI가 여러 번의 느린 단계 대신 단 한 번의 도약으로 이미지를 생성해야 하는 까다로운 '원스텝(one-step)' 설정에서도 이 다양한 접근 방식은 더 효과적이었습니다. 하지만 재미는 사진에서 끝나지 않았습니다. 연구팀은 이 방식을 생명의 기본 단위인 새로운 단백질을 설계하거나 3D 분자를 생성하는 것과 같은 과학적 과업에도 적용했습니다. 이 분야들에서도 AI가 더 다양한 내부 표현을 갖도록 장려하는 것이 더 나은 결과로 이어졌습니다.

이 논문은 AI의 학습을 가이드하기 위해 거대한 외부 모델이 필요하다는 생각에 반론을 제기합니다. 다른 방법들은 AI를 외부의 '전문가' 모델과 일치시키려 노력하지만, 저자들은 AI 자신의 내부 구성 요소들을 더 다양하게 만드는 것만으로도 최고 수준의 성능을 내기에 충분하다는 것을 발견했습니다. 또한 그들은 너무 많은 뇌의 부위를 외부 전문가와 일치시키려 하는 것이 도움이 되지 않으며, 오히려 성능을 해칠 수 있다는 점도 보여주었습니다. 핵심적인 교훈은 다양성이 보편적인 원리라는 것입니다. 당신이 고양이를 그리든 단백질을 접든, 동시에 여러 가지 다른 방식으로 생각하는 두뇌가 더 나은 것을 창조하는 두뇌입니다. 저자들은 이 접근 방식이 무거운 외부의 지도에 의존하지 않고도, 미래에 훨씬 더 강력하고 효율적인 AI를 여는 열쇠가 될 수 있다고 제안합니다.

기술 요약: DiverseDiT++

문제 정의

최근 Diffusion Transformer(DiT)의 발전은 그 확장성 덕분에 시각적 합성 분야에서 놀라운 진전을 이루었습니다. 그러나 DiT 내부의 표현 학습(representation learning)을 지배하는 근저의 메커니즘은 여전히 제대로 이해되지 않고 있습니다. REPA와 같은 최근 연구들은 내부 표현을 외부의 사전 학습된 인코더와 정렬함으로써 DiT 학습을 개선하려고 시도하지만, 이러한 방법들은 방대한 외부 자원에 의존하며 왜 그러한 방식이 효과적인지에 대해 충분히 설명하지 못합니다. 또한, 외부 가이드 없이 표현을 향상시키려는 기존의 접근 방식들(예: 분산 손실(dispersive losses))은 내부 표현이 어떻게 진화하는지에 대한 원리적인 이해가 부족합니다. 핵심적인 질문은 다음과 같습니다: DiT 모델은 어떻게 의미 있는 표현을 학습하는가, 그리고 왜 외부 정렬 기술이 효과적인가? 이러한 근본적인 이해의 결여는 더 효율적이고 원리적인 학습 패러다임의 개발을 저해합니다.

방법론

본 논문은 외부 가이드 모델에 의존하지 않고 DiT 내에서 다양한 표현 학습을 명시적으로 촉진하도록 설계된 프레임워크인 **DiverseDiT++**를 제안합니다. 이 방법론은 표현 역학(representation dynamics)에 대한 체계적인 분석을 기반으로 하며 두 가지 주요 구성 요소로 이루어집니다.

1. 체계적 분석 및 가중 다양성 점수 (WDS)

솔루션을 제안하기에 앞서, 저자들은 서로 다른 트랜스포머 블록 간의 유사성을 측정하기 위해 **Centered Kernel Alignment (CKA)**를 사용하여 DiT 표현 역학에 대한 체계적인 분석을 수행했습니다. 주요 관찰 결과는 다음과 같습니다:

블록 간 표현 불일치(discrepancy)는 학습 과정 중에 자연스럽게 증가합니다.
특정 블록을 사전 학습된 모델과 정렬하면 다른 블록들과의 비유사성이 증가합니다.
과도한 정렬(여러 블록 또는 여러 인코더 사용)은 반드시 성능을 향상시키는 것은 아니며, 전체적인 다양성을 해칠 수 있습니다.
다양성 패턴은 고유하며 노이즈 제거 타임스텝(denoising timesteps) 전반에 걸쳐 일관되게 나타납니다.

이러한 발견을 바탕으로, 저자들은 모든 트랜스포머 블록에 대한 쌍별(pairwise) CKA 유사도의 거리 가중치 보수(distance-weighted complement)로 정의되는 스칼라 지표인 **가중 다양성 점수 (Weighted Diversity Score, WDS)**를 도입했습니다:
$WDS = 1 - \frac{\sum_{i<j} |j-i| \cdot CKA(H_i, H_j)}{\sum_{i<j} |j-i|}$
저자들은 WDS와 log(FID) 사이에 강한 음의 상관관계( $r = -0.869$ )가 있음을 발견하였으며, 이를 통해 WDS가 모델 성능 및 생성 품질의 신뢰할 수 있는 지표임을 입증했습니다.

2. DiverseDiT++ 프레임워크

고성능 모델에서 관찰되는 다양성을 능동적으로 촉진하기 위해, DiverseDiT++는 두 가지 메커니즘을 도입합니다:

Long Residual Connections (장거리 잔차 연결): 균질한 입력으로 인해 발생하는 표현의 균질화(homogenization)를 방지하기 위해, 초기 레이어( $i$ )의 출력을 후기 레이어( $L-i$ )로 주입하는 long skip connection을 사용합니다. 이는 깊은 블록을 위한 입력 표현을 다양화하여 특징 재사용을 장려하고 붕괴(collapse)를 방지합니다.
Representation Diversity Loss (표현 다양성 손실): 서로 다른 블록의 특징 간 유사성에 벌점을 부과하는 명시적인 손실 함수를 도입합니다. 이는 세 가지 구성 요소로 이루어집니다:
1. Orthogonality Loss (직교성 손실): 블록별 평균 표현 간의 높은 코사인 유사성에 벌점을 부과합니다.
2. Mutual Information Loss (상호 정보량 손실): 블록 표현 간의 상호 정보량을 근사하고 최소화하여 통계적 독립성을 보장합니다.
3. Feature Dispersion Loss (특징 분산 손실): 특징 활성화의 분산을 최대화하여 다양한 채널 사용을 독려합니다.
전체 다양성 손실은 그래디언트 충돌 분석에서 유도된 **적응형 가중치 전략(adaptive weighting strategy)**에 의해 가중치가 조절됩니다. 이 가중치는 다양성이 포화될 때 그래디언트 충돌을 방지하고 모델이 후기 단계에서 미세한 디테일에 집중할 수 있도록 학습 과정 동안 클리핑(clipping) 및 감쇠(decay)됩니다.

주요 기여

메커니즘적 통찰: 본 논문은 블록 간의 표현 다양성이 DiT의 효과적인 학습을 위한 핵심 요소임을 체계적으로 밝혀낸 최초의 연구이며, 기존 정렬 기술의 성공(즉, 특화(specialization)를 강제함으로써 작동함)에 대한 메커니즘적 설명을 제공합니다.
WDS 지표: 규모, 방법 및 해상도 전반에 걸쳐 생성 품질(FID)과 강력하게 상관되는, 블록 간 다양성을 정량적으로 포착하는 원리적인 지표인 가중 다양성 점수(WDS)를 도입했습니다.
DiverseDiT++ 프레임워크: 외부 사전 학습된 인코더 없이 long residual connection과 적응형 다양성 손실을 통해 표현 학습을 강화하는 실용적이고 아키텍처에 구애받지 않는 프레임워크를 제안합니다.
일반화: 이미지 생성(ImageNet), 단백질 역전사(protein inverse folding), 3D 분자 생성 등 다양한 도메인에서의 검증을 완료했습니다.

실험 결과

저자들은 ImageNet (256×256 및 512×512), 단백질 역전사, 3D 분자 생성에 대해 광범위한 실험을 수행했습니다.

이미지 생성: DiverseDiT++는 다양한 백본(SiT, REPA)과 모델 규모(B, L, XL)에서 일관되게 성능을 향상시켰습니다.
- ImageNet 256×256에서 DiverseDiT++를 REPA-B에 적용했을 때 FID가 22.99에서 17.22(400k iterations)로 감소하여, 동일한 학습 비용을 가진 SiT-L 베이스라인(18.77)보다 우수한 성능을 보였습니다.
- 이 방법은 수렴을 가속화하여 200 에포크 만에 FID 1.49를 달ach하여, 1400 에포크 동안 학습된 SiT-XL/2(FID 2.06)를 능가했습니다.
- MeanFlow-XL/2를 사용하여 ImageNet 256×256의 one-step 생성에서 새로운 SOTA FID인 2.94를 달성했습니다.
과학적 도메인:
- 단백질 역전사: PDB 데이터셋에서 서열 회복률(sequence recovery rate)을 38.30%로 개선하고 scRMSD를 1.81로 낮추어 ProteinMPNN 및 REED와 같은 베이스라인보다 우수한 성능을 보였습니다.
- 3D 분자 생성: GEOM-Drug 및 QM9 데이터셋에서 SemlaFlow 및 REED와 비교하여 안정성을 높이고 에너지/변형(energy/strain)을 줄였습니다.
Ablation 연구: Long residual connection과 전체 다양성 손실(직교성, MI, 분산)이 최적의 성능을 위해 모두 필요함을 확인했습니다. 또한 이 방법이 DispLoss 및 SRA와 같은 기존 기술과 상보적임을 보여주었습니다.

의의 및 주장

본 논문은 블록 간의 표현 다양성이 시각적 합성을 넘어 과학적 도메인까지 아우르는 효과적인 확산 모델 학습을 위한 보편적인 원리라고 주장합니다. WDS를 통해 이 다양성을 정량화함으로써, 저자들은 DiT 학습 역학을 이해하기 위한 진단 도구를 제공합니다. 제안된 DiverseDiT++ 프레임워크는 외부 정렬 모델의 계산 오버헤드 없이도 내부 다양성을 명시적으로 장려함으로써 성능을 높이는 실용적인 전략을 제공합니다. 이 연구는 내부 다양성을 촉진하는 것이 다양한 데이터 모달리티와 아키텍처 전반에서 생성 모델링을 향상시킬 수 있는 일반화 가능한 귀납적 편향(inductive bias)임을 시사합니다.

저자들은 파라미터 수의 약간의 증가(잔차 연결로 인한)와 새로운 태스크를 위한 적응형 가중치 임계값의 경험적 결정 등 한계점을 언급하며, 이에 대한 추가적인 체계적 탐색이 필요할 수 있다고 밝혔습니다. 또한 텍스트-투-이미지(text-to-image) 및 이미지 편집으로의 일반화 가능성은 향후 과제로 남아 있음을 인정했습니다.