DiverseDiT: Towards Diverse Representation Learning in Diffusion Transformers

이 논문은 Diffusion Transformer(DiT) 의 내부 표현 다양성이 학습 성능에 결정적임을 규명하고, 이를 촉진하기 위해 긴 잔차 연결과 표현 다양성 손실 함수를 도입한 새로운 프레임워크인 DiverseDiT 를 제안하여 다양한 백본과 조건에서 성능 향상과 수렴 가속화를 입증합니다.

핵심

🎨 DiverseDiT: 그림을 그리는 AI 에게 "다양한 생각"을 심어주다

이 논문은 최신 이미지 생성 AI(확산 모델) 가 더 좋은 그림을 그리기 위해 필요한 핵심 비결을 발견하고, 이를 해결하는 새로운 방법을 제안한 연구입니다.

핵심 아이디어를 한 마디로 요약하면:

"AI 가 그림을 그릴 때, 각 단계마다 똑같은 생각만 하지 않고, 서로 다른 관점과 아이디어를 나누게 하면 훨씬 더 훌륭한 그림이 나온다."

이제 이 복잡한 개념을 일상적인 비유로 쉽게 설명해 드릴겠습니다.

---

1. 문제점: "똑같은 생각만 하는 팀" (기존 방식의 한계)

그림을 그리는 AI 를 상상해 보세요. 이 AI 는 여러 층 (Layer) 으로 이루어진 거대한 공장처럼 작동합니다.

• 기존 방식: 공장의 1 층에서 재료를 다듬고, 2 층에서 색을 입히고, 3 층에서 디테일을 추가합니다.
• 문제: 그런데 이상하게도, 이 공장 각 층들이 서로 똑같은 생각을 하고 있었습니다. 1 층이 "고양이 귀를 그려야지"라고 생각하면, 2 층, 3 층, 4 층도 모두 "고양이 귀"만 생각하며 똑같은 작업을 반복했습니다.
• 결과: AI 는 그림을 그리다가 지루해지고 (학습 효율 저하), 결과물도 평범하거나 뻔한 모습만 만들어냈습니다. 이를 전문 용어로 **'표현의 동질화 (Homogenization)'**라고 합니다.

2. 기존 해결책: "외부 전문가의 지시" (REPA 등)

이 문제를 해결하기 위해 이전 연구자들은 **외부 전문가 (사전 훈련된 다른 AI)**를 고용했습니다.

• 방식: "1 층은 전문가 A 의 지시를 듣고, 2 층은 전문가 B 의 지시를 들어라"라고 시켰습니다.
• 단점: 전문가를 고용하려면 엄청난 비용 (컴퓨팅 자원) 이 들고, 전문가가 너무 많으면 오히려 공장 내부가 혼란스러워져서 더 나쁜 결과가 나오기도 했습니다.

3. DiverseDiT 의 해결책: "내부 팀워크 강화"

이 논문 (DiverseDiT) 은 외부 전문가 없이도 공장 내부에서 자연스럽게 다양한 생각이 나오게 하는 두 가지 간단한 방법을 제안했습니다.

🌟 방법 1: "오래된 아이디어도 다시 꺼내보자" (긴 잔여 연결, Long Residual Connections)

• 비유: 공장의 1 층에서 만든 아이디어를 10 층까지 바로 전달하는 게 아니라, 1 층의 아이디어를 5 층, 10 층 등 먼 곳으로 직접 연결해 주는 통로를 만듭니다.
• 효과: 10 층이 9 층의 결과만 보는 게 아니라, 1 층의 초기 아이디어도 함께 보고 "아, 처음엔 이런 생각이 있었구나!"라고 생각하게 됩니다. 이렇게 하면 각 층이 서로 다른 정보를 바탕으로 작업하게 되어 생각이 다양해집니다.

🌟 방법 2: "서로 다른 생각을 하라" (다양성 손실 함수, Diversity Loss)

• 비유: 공장 관리자가 각 층에 **"너희는 서로 너무 비슷하게 생각하면 안 돼!"**라고 경고하는 규칙을 만듭니다.
• 방식: 1 층과 2 층의 생각이 너무 비슷하면 (유사도가 높으면) 벌점을 주고, 서로 다른 생각을 하면 점수를 줍니다.
• 효과: 각 층은 서로 다른 역할을 하려고 노력합니다. 1 층은 '형체'를 잡고, 2 층은 '색감'을, 3 층은 '질감'을 담당하는 식으로 전문가처럼 분업하게 됩니다.

---

4. 결과: 왜 이 방법이 좋은가요?

이 방법을 적용한 실험 결과, 놀라운 변화가 일어났습니다.

1. 더 빠른 학습: 외부 전문가를 고용할 필요 없이, AI 스스로가 더 빠르게 배우고 발전했습니다. (학습 시간 단축)
2. 더 높은 품질: 생성된 그림의 디테일이 훨씬 풍부해지고, 예술적인 완성도가 높아졌습니다.
3. 범용성: 작은 AI 모델부터 거대한 AI 모델까지, 모든 크기의 모델에 적용되어 좋은 성과를 냈습니다.
4. 한 걸음 생성도 가능: 보통 그림을 그리려면 여러 번 반복해야 하지만, 이 방법을 쓰면 한 번에 고품질 그림을 그릴 수도 있게 되었습니다.

📝 요약: 이 논문이 우리에게 주는 교훈

이 연구는 **"AI 가 좋은 것을 만들기 위해서는, 내부 구성원들이 서로 다른 관점을 가지고 협력해야 한다"**는 사실을 증명했습니다.

• 과거: "외부 전문가의 지시를 받아라." (비싸고 복잡함)
• 현재 (DiverseDiT): "내부 팀원들이 서로 다른 생각을 하도록 유도하라." (간단하고 효율적)

마치 훌륭한 오케스트라가 지휘자의 지시 없이도 각 악기들이 서로 다른 소리를 내며 조화로운 음악을 만들어내듯, DiverseDiT 는 AI 가 스스로 다양한 아이디어를 만들어내어 더 아름다운 그림을 그리게 해주는 혁신적인 방법입니다.

기술 요약

1. 문제 정의 (Problem)

최근 확산 모델 (Diffusion Models), 특히 **Diffusion Transformers (DiT)**는 시각적 합성 분야에서 뛰어난 확장성과 성능을 보여주고 있습니다. 기존 연구들 (예: REPA) 은 사전 훈련된 외부 인코더를 사용하여 DiT 의 내부 표현 (internal representations) 을 정렬 (alignment) 함으로써 성능을 향상시켰습니다.

그러나 DiT 가 어떻게 의미 있는 내부 표현을 학습하는지에 대한 근본적인 메커니즘은 여전히 불명확합니다. 기존 방법론들은 다음과 같은 한계가 있습니다:

• 외부 의존성: REPA 와 같은 방법은 대규모 사전 훈련된 모델 (Foundation Models) 에 의존하여 막대한 계산 자원을 소모합니다.
• 표현의 동질화 (Homogenization): DiT 의 여러 블록 (layers) 간 입력이 균일하고, 블록 간 표현이 서로 유사해져서 (representation collapse) 학습 효율이 떨어질 수 있습니다.
• 기작의 불명확성: 외부 정렬이 왜 효과적인지, 그리고 DiT 내부에서 표현이 어떻게 진화하는지에 대한 체계적인 분석이 부족했습니다.

2. 방법론 (Methodology)

저자들은 DiT 의 표현 학습 역학을 체계적으로 분석한 후, **표현의 다양성 (Representation Diversity)**이 학습의 핵심 요소임을 발견하고 이를 명시적으로 증진시키는 새로운 프레임워크 DiverseDiT를 제안합니다.

2.1. 핵심 통찰 (Key Insights from Analysis)

저자들은 CKA (Centered Kernel Alignment) 를 사용하여 DiT 의 내부 표현을 분석한 결과, 다음과 같은 사실을 발견했습니다:

1. 학습 진행에 따른 다양성 증가: 훈련이 진행됨에 따라 서로 다른 블록 간의 표현 차이 (discrepancy) 가 자연스럽게 증가합니다.
2. 외부 정렬의 효과: 특정 블록을 외부 모델에 정렬하면 해당 블록과 다른 블록 간의 차이가 커지며, 이는 블록의 전문화 (specialization) 를 촉진합니다.
3. 과도한 정렬의 역효과: 너무 많은 블록을 외부 모델에 정렬하거나 여러 인코더를 사용하면 오히려 전체적인 표현 다양성이 감소하여 성능이 저하될 수 있습니다.

2.2. DiverseDiT 의 구성 요소

이러한 통찰을 바탕으로, 외부 모델 없이 내부 메커니즘만으로 표현 다양성을 증진시키는 두 가지 핵심 기법을 도입했습니다.

1. 긴 잔차 연결 (Long Residual Connections):

   • 기존 DiT 는 각 블록의 입력이 이전 레이어의 출력에만 의존하여 동질화되는 경향이 있습니다.
   • DiverseDiT 는 초기 레이어의 출력을 후속 레이어에 직접 주입하는 **긴 잔차 연결 (Long-range residual connections)**을 도입합니다.
   • 이는 각 블록으로 들어가는 입력의 다양성을 높이고, 표현의 동질화를 방지하여 더 다양한 특징을 학습하도록 유도합니다.

2. 표현 다양성 손실 함수 (Representation Diversity Loss, $L_{div}$):

   • 서로 다른 블록 간의 특징이 서로 유사하지 않도록 명시적으로 패널티를 부여하는 손실 함수를 설계했습니다. 이는 세 가지 하위 손실로 구성됩니다:
     • 직교성 손실 ($L_{orth}$): 블록 간 평균 특징 벡터의 코사인 유사도를 최소화하여 직교성을 유도합니다.
     • 상호 정보 최소화 손실 ($L_{MI}$): 블록 간 통계적 독립성을 보장하기 위해 정규화된 특징 벡터 간의 상관관계를 최소화합니다.
     • 특질 분산 손실 ($L_{disp}$): 특징 활성화의 분산을 최대화하여 채널 사용의 다양성을 증진시킵니다.
   • 전체 손실은 $L_{div} = \lambda_{orth}L_{orth} + \lambda_{MI}L_{MI} + \lambda_{disp}L_{disp}$로 정의되며, 학습 안정성을 위해 손실 값에 따라 가중치를 적응적으로 조절합니다.

3. 주요 기여 (Key Contributions)

1. DiT 표현 학습 역학에 대한 체계적 분석: DiT 내부에서 블록 간 표현 다양성이 학습 성능에 결정적인 역할을 한다는 것을 최초로 규명하고, 기존 외부 정렬 방법 (REPA 등) 의 효과를 설명하는 새로운 관점을 제시했습니다.
2. DiverseDiT 프레임워크 제안: 외부 모델 없이도 긴 잔차 연결과 표현 다양성 손실을 통해 효율적이고 강력한 표현 학습을 가능하게 하는 새로운 아키텍처를 제안했습니다.
3. 광범위한 실험 검증: ImageNet 256x256 및 512x512 데이터셋에서 다양한 모델 규모 (SiT, REPA, MeanFlow 등) 와 설정 (다단계, 1 단계 생성) 에서 일관된 성능 향상을 입증했습니다.

4. 실험 결과 (Results)

• 성능 향상: ImageNet 256x256 에서 SiT-B 모델에 DiverseDiT 를 적용했을 때, FID 점수가 36.80 에서 28.05로 크게 개선되었습니다. REPA-B 모델의 경우 22.99 에서 17.29로 향상되었으며, 이는 더 큰 규모의 모델 (REPA-L) 보다 더 좋은 성능을 내기도 했습니다.
• 수렴 가속: DiverseDiT 는 기존 방법론보다 훨씬 적은 훈련 에포크 (예: 200 에포크) 에서 SOTA 모델 (수천 에포크 훈련) 과 견줄 만한 성능을 달성하여 수렴 속도를 획기적으로 높였습니다.
• 1 단계 생성 (One-step Generation): MeanFlow 기반의 1 단계 생성 모델에서도 FID 2.99 를 기록하여 새로운 SOTA 를 달성했습니다.
• 상호 보완성: DiverseDiT 는 DispLoss 나 SRA 와 같은 기존 내부 정렬 기법과도 결합되어 추가적인 성능 향상을 제공하며, 외부 모델 의존성 없이도 REPA 보다 우수한 결과를 보여줍니다.

5. 의의 및 중요성 (Significance)

• 자원 효율성: 거대하고 비싼 외부 사전 훈련 모델 (Foundation Models) 에 대한 의존성을 제거하면서도, 혹은 그보다 더 나은 성능을 달성함으로써 생성 모델의 학습 비용을 대폭 절감할 수 있습니다.
• 이론적 기여: 확산 모델이 어떻게 내부적으로 표현을 학습하는지에 대한 '블록 간 다양성'이라는 핵심 원리를 규명함으로써, 향후 더 효율적이고 원리 기반 (principled) 인 학습 패러다임을 설계하는 데 중요한 통찰을 제공합니다.
• 실용성: 다양한 모델 아키텍처와 규모에 적용 가능하며, 외부 가이드 없이도 내부 메커니즘만으로 표현의 질을 높일 수 있어 실제 적용 가능성이 매우 높습니다.

요약하자면, DiverseDiT는 확산 트랜스포머의 내부 표현이 서로 다른 블록 간에 얼마나 다양해야 하는지 규명하고, 이를 위한 간단한 구조적 변경과 손실 함수를 통해 외부 의존성 없이도 최상의 생성 성능을 달성하는 혁신적인 방법론입니다.