새로운 딥러닝 아키텍처가 발표될 때마다 연구자들을 괴롭히는 것이 하나 있습니다. 바로 '하이퍼파라미터 튜닝'이죠. 성능을 쥐어짜내기 위해 학습률을 깎고, 웜업(Warmup) 스케줄을 조정하며 며칠 밤낮을 새우는 것이 일상입니다.
하지만 DiT(Diffusion Transformer) 논문의 훈련(Training) 파트를 읽다 보면, 저자들의 뼈 있는 '자랑'에 놀라게 됩니다. "우리 모델은 이것저것 복잡한 기법을 안 써도 알아서 학습이 엄청나게 잘 되더라!"라는 점 때문입니다.
오늘은 DiT가 어떻게 그토록 안정적인 학습을 이뤄냈는지 그 실험 설정의 비밀을 파헤쳐 보고, 생성 모델의 퀄리티를 한 단계 높여주는 실무 비법인 EMA(지수 이동 평균) 기법까지 깊이 있게 알아보겠습니다.
1. DiT의 훈련 설정: 완벽한 뼈대가 주는 자신감
DiT 논문의 실험 설정(Experimental Setup)을 살펴보면, 이 모델이 얼마나 튼튼하게 설계되었는지 알 수 있습니다.
- 데이터셋 & 해상도: 경쟁이 치열한 ImageNet 데이터셋을 사용해 256x256 및 512x512 해상도로 훈련했습니다.
- 옵티마이저: AdamW를 사용했으며, 학습률은 로 일정하게 유지했고, 배치 크기는 256으로 설정했습니다.
- 데이터 증강(Augmentation): 복잡한 기법 없이 오직 '수평 뒤집기(Horizontal Flips)' 하나만 사용했습니다.
💡 여기서 주목해야 할 가장 놀라운 점!
보통 트랜스포머(Transformer) 모델을 훈련할 때는 손실 값 급증(Loss Spikes)을 막기 위해 '학습률 웜업'이나 '가중치 감소(Weight Decay)' 같은 정규화 기법이 필수적입니다. 하지만 DiT는 이런 기법들을 전혀 사용하지 않았습니다.
심지어 모델의 크기(S부터 XL까지)나 패치 크기가 달라져도, 기존 U-Net 기반 모델(ADM)이 쓰던 설정값을 튜닝 없이 그대로 복사해서 썼습니다. 앞서 적용한 'adaLN-Zero' 같은 초기화 기법(마지막 선형 계층을 0으로 초기화) 덕분에 모델 구조 자체가 너무나도 안정적이어서, 무식하게 밀어붙여도 성공적인 학습이 가능했던 것입니다.
2. 생성 모델의 디테일을 살리는 마법, EMA (지수 이동 평균)
저자들이 훈련 과정에서 일반적인 관행을 따라 필수적으로 유지한 것이 하나 있습니다. 바로 EMA(Exponential Moving Average)입니다. 논문에서는 0.9999의 감소율(Decay rate)을 적용해 모든 결과를 보고했습니다.
EMA란 무엇일까요?
학습 과정에서 가중치는 매 스텝마다 조금씩 진동(Oscillation)하며 업데이트됩니다. EMA는 최근 가중치들의 평균을 내어 아주 부드럽고 묵직하게 가중치를 갱신하는 기법입니다. 수식으로는 다음과 같이 표현됩니다.
이 기법을 도입하면 모델이 놓쳤던 미세한 마진이나 선명한 그루브 등 3D 및 2D 생성 모델의 디테일이 훨씬 매끄럽고 안정적으로 도출됩니다. 현업에서 퀄리티를 쥐어짜내는 '숨겨진 비법'이기도 하죠.
3. 실무자를 위한 파인튜닝 + EMA 가이드
그렇다면 우리가 직접 9.5M 크기의 사전 학습 모델을 가져와 파인튜닝(Fine-tuning)을 할 때, EMA를 어떻게 설정해야 할까요? 붕괴율(Decay Rate, )에 따라 요구되는 최소 학습 스텝 수가 다릅니다.
- (빠른 반영): 최근 약 100번의 스텝을 평균 냅니다. 최소 500 ~ 1,000 스텝만 업데이트되어도 평활화(Smoothing) 효과를 볼 수 있습니다. 데이터 수가 적고 빠른 실험이 필요할 때 적합합니다.
- (표준 추천값): 최근 약 1,000번의 스텝을 평균 냅니다. 의료용 3D 비전이나 생성 모델에서 가장 널리 쓰입니다. 최소 3,000 ~ 5,000 스텝 이상 돌아가야 진가를 발휘합니다.
- (논문 세팅, 매우 보수적): 스크래치(Scratch)부터 학습할 때 주로 쓰며, 10,000~20,000 스텝 이상이 필요합니다. 규모가 작은 파인튜닝 환경에서는 너무 느려 추천하지 않습니다.
🚨 파인튜닝 시 치명적인 주의사항 (Gotcha)
파인튜닝 시 스텝 0에서 EMA 가중치를 랜덤이나 0으로 초기화하면 절대 안 됩니다! 반드시 사전 학습이 완료된 원본 모델의 가중치(Checkpoint)를 EMA 모델의 초기 가중치로 똑같이 복사(Deepcopy) 해 놓고 시작해야 합니다. 그래야 파인튜닝 시작 직후부터 튼튼한 뼈대 위에서 안정적으로 디테일을 깎아나갈 수 있습니다.
마무리하며
DiT는 강력한 아키텍처 설계가 훈련 과정을 얼마나 단순하고 우아하게 만들어주는지 증명한 훌륭한 사례입니다. 여기에 EMA라는 디테일 장인을 결합하면, 여러분의 파인튜닝 모델도 한 차원 높은 결과물을 만들어낼 수 있을 것입니다.
