μP를 통한 확산 변환기의 효율적인 확장

확산 변환기(Diffusion Transformers)는 시각 생성 모델의 기반이 되었지만, 대규모에서 하이퍼파라미터(HP) 튜닝 비용이 높아 확장성이 제한적입니다. 최근에, 일반 변환기(vanilla Transformers)를 위한 최대 업데이트 매개변수(Maximal Update Parametrization, muP)가 제안되어, 작은 언어 모델에서 큰 언어 모델로 안정적인 HP 전송을 가능하게 하고 튜닝 비용을 크게 줄였습니다. 그러나 일반 변환기와 구조적으로 및 목적상 차이가 있는 확산 변환기에 muP가 적용될 수 있는지는 여전히 불분명합니다. 본 연구에서는 표준 muP를 확산 변환기에 일반화하고 대규모 실험을 통해 그 효과성을 검증하였습니다. 첫째, DiT, U-ViT, PixArt-alpha, MMDiT 등 주요 확산 변환기의 muP가 일반 변환기의 muP와 일치함을 엄밀히 증명하여 기존의 muP 방법론을 직접 적용할 수 있게 하였습니다. 이 결과를 활용하여, DiT-muP의 강건한 HP 전송 가능성을 체계적으로 입증하였습니다. 특히, 학습률이 전송된 DiT-XL-2-muP는 원래의 DiT-XL-2보다 2.9배 더 빠른 수렴성을 보였습니다. 마지막으로, PixArt-alpha를 0.04B에서 0.61B로, MMDiT를 0.18B에서 18B로 확장하여 텍스트-이미지 생성에서 muP의 효과성을 검증하였습니다. 두 경우 모두 muP 아래에서 동작하는 모델들이 각각의 베이스라인보다 우수한 성능을 보였으며, PixArt-alpha는 단 한 번의 학습 과정 비용의 5.5%만으로, MMDiT-18B는 인간 전문가에 의한 소비량의 3%만으로 적은 튜닝 비용으로 성능을 발휘하였습니다. 이러한 결과들은 muP가 확산 변환기를 확장하는 원칙적이고 효율적인 프레임워크임을 확인해주었습니다.