UniGP: Taming Diffusion Transformer for Prior-Preserved Unified Generation and Perception

기술 요약: UNIGP – 사전 학습된 생성 프라이어를 보존하는 통합 생성 및 인식을 위한 디퓨전 트랜스포머 제어

문제 정의

최근의 디퓨전 모델 연구는 제어 가능한 이미지 생성(예: 깊이 또는 에지를 가이드로 하는 Text-to-Image)과 밀집 예측 작업(예: 단안 깊이 및 표면 법선 추정) 모두에서 눈부신 성과를 거두었습니다. 그러나 기존의 접근 방식들은 대개 이들을 별개의 영역으로 취급합니다. 제어 가능한 생성 방법(ControlNet 등)은 생성 작업에 국한되며 정밀한 제어를 위해 방대한 데이터셋을 요구하는 경우가 많고, 밀집 예측 방법(Marigold 또는 회귀 기반 모델 등)은 생성 백본으로부터 파인튜닝되면서 모델의 다재다능한 생성 프라이어를 상실하는 "파괴적 망각(catastrophic forgetting)" 현상을 겪곤 합니다.

현재의 통합 시도들(예: JointNet, UniCon)은 생성과 인식 사이의 상호 보완적 관계를 충분히 활용하지 못하는 경우가 많습니다. 이들은 전체 백본을 복제함으로써 막대한 계산 오버헤드를 초래하거나, 백본을 인식 작업을 위해 직접 파인튜닝함으로써 생성 능력을 저해합니다. 핵심 과제는 복잡한 작업별 아키텍처나 손실 함수 없이도, 사전 학습된 생성 프라이어를 보존하면서 정확한 밀집 예측과 제어 가능한 생성을 가능하게 하는 단일 프레임워크 내에 이러한 이질적인 작업들을 통합하는 것입니다.

방법론: UNIGP 프레임워크

저자들은 MMDiT(Multi-Modal Diffusion Transformer) 아키텍처를 기반으로 한 UNIGP 프레임워크를 제안합니다. 이 방법은 두 가지 주요 혁신인 분리된 아키텍처 브랜치와 통합된 학습 전략으로 구성됩니다.

1. 분리된 통합 생성 및 인식 (DUGP)

오컴의 면도날 원칙에 따라, UNIGP는 전체 MMDiT 백본을 복제하는 대신 DUGP 브랜치를 도입하여 이를 피합니다. DUGP 브랜치는 사전 학습된 MMDiT의 이미지 브랜치($\mathcal{F}_i$)만을 복사합니다.

• 아키텍처: DUGP 브랜치는 이미지 브랜치의 파라미터를 복사하여 초기화됩니다. 이는 Condition Patcher, Stacked Control Layers, 그리고 Stacked Perception Layers로 나뉩니다.
• 프라이어 보존 (그래디언트 격리): 인식 레이어에서 가산 연결(additive connection)을 우회하는 것은 결정적인 설계 선택입니다. 제어 레이어에서는 DUGP 브랜치의 출력이 제로 초기화된 선형 레이어($\mathcal{L}_{zero}$)를 통해 백본의 출력에 더해져 조건이 생성을 가이드할 수 있게 합니다. 그러나 인식 레이어에서는 이 가산 경로가 차단됩니다. 이는 밀집 예측 손실로부터 오는 그래디언트가 오직 DUGP 브랜치를 통해서만 흐르게 하여, 백본의 사전 학습된 가중치를 변경하는 것을 방지합니다. 이는 생성 프라이어를 효과적으로 격리(quarantine)하는 동시에, DUGP 브랜치가 비침습적인 "독자(reader)"로서 인식을 학습할 수 있도록 합니다.
• 다중 조건 처리: 여러 조건이 필요한 시나리오의 경우, 추론 효율성을 유지하기 위해 DUGP 컨디션 패처를 각 조건마다 복제하고 상호작용은 시퀀스 결합(sequence concatenation)을 통해 처리합니다.

2. 통합 데이터셋 및 학습 전략

생성과 인식 사이의 간극을 메우기 위해, UNIGP는 필터링된 생성 데이터와 합성된 인식 데이터를 결적으로 결합한 이질적인 데이터셋을 사용하여 통합 학습 전략을 채택합니다.

• 이진 손실 가중치(Binary Loss Weighting): 복잡한 손실 균형 대신, UNIGP는 단순한 이진 라우팅 메커니즘을 사용합니다. 혼합 배치 내의 각 샘플에 대해, 데이터 소스에 따라 손실 가중치($\lambda_g, \lambda_p$)를 1 또는 0으로 설정합니다. 샘플이 생성 데이터셋에서 온 것이라면 생성 손실이 활성화되고, 인식 데이터셋에서 온 것이라면 인식 손실이 활성화됩니다.
• 목표: 네트워크는 RGB, 깊이, 법선 잠재 변수(latents)에 대한 속도장(velocity field)을 동시에 예측하도록 학습되어, 단일 순전파 내에서 생성 품질과 기하학적 정확도를 모두 최적화합니다.

주요 기여

1. UNIGP 프레임워크: 제어 가능한 생성과 밀집 예측을 통합하는 통합 디퓨전 트랜스포머 모델입니다. 이는 오컴의 면도를 준수하여 계산 오버헤드를 최소화하면서 MMDiT의 이미지 브랜치만을 복사하는 DUGP 브랜치를 활용하여, 백본의 다재다능한 프라이어를 보존합니다.
2. 통합 학습 전략: 단순한 이진 손실 가중치 방식을 사용하여 이질적인 작업을 공동으로 학습하는 새로운 데이터셋 및 학습 접근법을 제시하며, 이를 통해 작업별 아키텍처 설계나 복잡한 손실 균형의 필요성을 제거했습니다.
3. 입증된 시너지: 논문은 공동 학습이 상호 보완적인 이점을 제공한다는 증거를 제시합니다: 생성 프라이어는 지각적 세부 사항(미세한 질감 포착)을 풍부하게 하고, 인식 학습은 생성 시 엄격한 공간적 정렬을 강제합니다.

실험 결과

광범적인 실험을 통해 UNIGP가 이전의 통합 접근 방식들을 능가하며 전문화된 "전문가(expert)" 방법들과 대등한 성능을 보임을 입증했습니다.

• 제어 가능한 생성: UNIGP는 다양한 조건(깊이, 법선, 스케치, 포즈, Canny)에 대해 단일 제어(ControlNet) 및 다중 제어(UniControlNet) 모델보다 FID 및 CLIP 점수에서 우수한 성능을 보였습니다. 특히, 깊이-이미지 및 법선-이미지 작업에서 가장 낮은 RMSE를 달성하여 탁월한 구조적 정렬을 보여주었습니다.
• 밀집 예측: 제로샷 깊이 및 표면 법선 추정에서, UNIGP는 훨씬 적은 이미지(DepthAnything V2의 62.6M 대비 59K)로 학습되었음에도 불구하고 최첨단 회귀 모델(예: DepthAnything V2, DSINE)과 대등한 성능을 달-성했습니다. 또한 다른 생성 기반 인식 방법(Marigold, GeoWizard) 및 통합 모델(JointNet, UniCon)보다 크게 앞섭러 있습니다.
• 결합 생성: UNIGP는 텍스트 프롬프트로부터 RGB, 깊이, 표면 법선 맵을 동시에 생성하는 기능을 지원하며, 이미지 품질(FID)과 기하학적 일관성(GD/GN 지표) 모두에서 LDM3D 및 JointNet과 같은 베이스라인을 능가합니다.
• 절제 연구(Ablation Studies):
  • 시너지: 인식 학습을 제거하면 생성된 이미지의 구조적 정렬이 저하(높은 RMSE)되는 반면, 생성 학습을 제거하면 미세한 디테일이 부족한 지나치게 매끄러운 깊이 맵이 생성됩니다.
  • 설계 선택: DUGP를 전체 백본 복제(JointNet 스타일) 또는 전체 파인튜닝(Marigold 스타일)으로 대체할 경우 성능이 현저히 떨어지며, 이는 분리된 그래디언트 흐름이 성공의 핵심임을 확인시켜 줍니다.
  • 레이어 균형: 적층된 제어 레이어와 인식 레이어의 동일한 할당(12:12)이 생성 제어와 인식 정확도 사이의 최적의 균형을 제공합니다.

의의 및 주장

논문은 UNIGP가 디퓨전 기반 합성 및 밀집 기하 예측이라는 분리된 분야를 성공적으로 연결했다고 주장합니다. 이 연구의 주요 의의는 생성 프라이어와 인식 제약이 서로 배타적인 것이 아니라 상호 보완적이라는 것을 입증했다는 데 있습니다.

• 프라이어 보존: 그래디언트 흐름을 격리함으로써, UNIGP는 특정 인식 작업 파인튜닝 시 흔히 상실되는 사전 학습된 MMDiT 백본의 고품질 생성 능력을 유지합니다.
• 효율성: 이 방법은 전체 백본을 복제하는 것에 비해 상대적으로 적은 파라미터 증가(2B에서 3B로)를 통해 최첨단 결과를 달성하며, 이는 계산적으로 효율적인 통합 솔루션임을 보여줍니다.
• 다재다능함: 본 프레임워크는 텍스트-이미지, 깊이/법선 추정, 모든 조건 생성(any-condition generation)을 포함한 광범위한 작업을 단일 모델 내에서 지원하여, 여러 전문 모델을 유지하는 것보다 더 다재다능한 대안을 제공합니다.

저자들은 UNIGP가 기본 모델에 비해 추가적인 계산 오버헤드를 도입하여 추론 시간을 약 5초에서 9.5초로 증가시킨다는 점을 인정하지만, 여러 전문화된 아키텍처를 하나의 고도로 다재다능한 통합 모델로 대체할 수 있는 능력에 의해 이 트레이드오프가 정당화된다고 설명합니다.