버려지던 투영 헤드를 재활용해 확산 트랜스포머 학습을 가속한다

무엇인가

이 논문은 확산 트랜스포머(diffusion transformer) 학습을 빠르게 만드는 표현 정렬(representation alignment) 기법의 사각지대를 다룬다. REPA는 사전학습된 비전 인코더(DINOv2 등)의 특징을 중간 레이어에 정렬 손실로 감독해 학습을 가속한다. 이때 정렬 손실의 그래디언트는 앞쪽(early) 레이어로 역전파되지만, 뒤쪽(later) 레이어는 velocity 예측 손실만으로 학습되기 때문에 그 의미론에 잘 접근한다는 보장이 없다. 저자들은 CKNNA로 정렬도를 측정해, 정렬에 쓰이는 MLP 투영 헤드 g_φ의 출력이 트랜스포머의 어느 레이어보다도 비전 인코더와 강하게 정렬되어 있음을 보인다. 그런데도 기존 방법들은 학습이 끝나면 g_φ를 버린다. 저자들은 이것이 학습된 가중치와 의미론을 낭비하는 것이라고 주장한다.

어떻게 동작하나

제안 방법은 RECON(Representation Conditioning)이다. 핵심은 g_φ를 버리지 않고 남겨두는 것이다. 학습 시에는 깨끗한 이미지가 있으므로 비전 인코더 특징 y0를 파라미터 없는 풀링으로 k개의 시맨틱 토큰으로 만들고(기본값 k=1), 모델의 은닉 차원에 맞춘 뒤(기본값은 zero-padding이라 파라미터가 늘지 않는다. Small 모델처럼 인코더 차원이 더 클 때만 1M 미만의 선형 매핑을 쓴다) 패치 토큰과 이어 붙여 뒤쪽 블록에 넣는다. 이는 언어모델의 teacher forcing과 같은 방식으로, 아직 학습 중이라 불안정한 g_φ 출력 대신 안정적인 정답 특징을 조건으로 준다. 추론 시에는 깨끗한 이미지가 없으므로 g_φ(h)의 출력을 풀링해 같은 자리에 넣는다. 뒤쪽 블록들은 패치 토큰과 시맨틱 토큰을 함께 처리하고, 패치 토큰은 임의 깊이에서 시맨틱 토큰에 어텐션한다. 시맨틱 토큰은 마지막 블록 전에 제거된다. 학습과 추론의 분포 차이를 줄이기 위해 시맨틱 토큰은 t < t_max = 0.5인 타임스텝에서만 주입한다. 손실 함수는 바뀌지 않고 파라미터도 추가되지 않는다.

무엇과 다른가

실험은 ImageNet 256×256, SD-VAE 잠재, SiT-B/2·L/2·XL/2, DINOv2-B, 정렬 깊이 8, 배치 크기 256에서 이뤄졌다. 클래스 조건부 생성에서 400K 스텝에 gFID 3.25를 기록해 REG(3.40), REPA-E(4.07)를 앞섰고, REPA는 같은 표에서 4M 스텝에 5.90이었다. 학습을 늘리면 1M 스텝 2.54, 2.4M 스텝 2.21까지 내려간다. 수렴 속도는 SiT 대비 63배, REPA 대비 25배 빠르다고 보고한다. Autoguidance를 적용한 시스템 수준 비교에서는 80 에폭에 2.09로, 1600 에폭 학습된 MaskDiT(2.28)와 1400 에폭의 DiT(2.27)를 앞선다. 480 에폭에는 1.49로 1400 에폭의 SiT(2.06)를 앞서고 800 에폭 REPA(1.42)와 비슷해진다.

어떻게 쓰나

텍스트-투-이미지에서도 개선이 나타난다. MS-COCO의 MMDiT 학습에서 FID는 8.26으로 REPA의 9.85보다 낮고 CLIP 점수는 0.2313으로 0.2273보다 높다(w=1). 더 큰 BLIP3-o 데이터셋의 NanoGen 설정에서는 FID 3.90 대 5.41, IS 24.31 대 22.40을 기록했다.

전제와 한계

생성 품질 외에 중간 특징의 의미론도 측정했다. 시맨틱 토큰을 받는 뒤쪽 레이어에서 평균 CKNNA가 21.2%, 선형 프로빙 정확도가 4.9% 올랐다. 또 샘플의 구조가 더 이른 시점에 정해져서, 네트워크 평가 횟수(NFE)를 100으로 줄여도 gFID 3.2, 50으로 줄여도 3.3을 유지한다. 같은 조건에서 REPA는 7.9에서 8.4, 10.8로 나빠진다. 추론 시 2.5배에서 5배 적은 평가로 비슷한 품질을 얻는다는 뜻이다.

일반화 실험도 있다. 모델 크기 B/L/XL 모두에서 REPA보다 좋았고, 20 에폭 기준 SiT-L/2에 RECON을 붙인 gFID 11.24가 REPA를 붙인 SiT-XL/2의 18.88보다 좋았다. VAE를 바꿔도 SD-VAE 7.84→3.25, InVAE 7.41→4.78, VA-VAE 4.95→4.02로 일관되게 개선됐다(VA-VAE는 이미 REPA 손실로 학습돼 개선폭이 작다). I-JEPA-H, DINOv2-B, MoCo-v3-B, SigLIP2-B, WebSSL-DINO-1B 다섯 인코더에서도 모두 향상됐다. 절제 실험에서는 t_max=0.5, k=1, zero-padding, 정렬 깊이 8을 기본값으로 골랐다.

개발자 관점에서 이 방법의 매력은 비용 구조다. 새 파라미터가 사실상 늘지 않고 손실 함수도 그대로 두면서, 이미 학습 파이프라인에 넣어둔 REPA류 정렬 헤드를 추론 경로에 다시 연결하기만 하면 된다. 확산 모델 학습 예산이나 추론 스텝 수가 병목인 팀이라면, 정렬 헤드를 체크포인트에서 삭제하는 관행을 재검토할 근거가 된다. 다만 조건부 생성에서 classifier-free guidance의 이득이 RECON에서는 줄어든다고 저자들이 밝히고 있어, guidance 전략은 별도로 튜닝해야 한다.

한계는 명확하다. RECON은 REPA식 정렬 손실로 이미 학습된 투영 헤드를 재사용하므로, 그런 손실 없이 학습된 모델에는 적용할 수 없다. 저자들도 이것이 정렬 자체의 필요성을 없애는 것이 아니라 헤드를 버리는 낭비를 없애는 것이라고 말한다. 또한 시맨틱 토큰 수 k를 늘리면 g_φ 예측 오차가 증폭되어 생성 품질이 나빠질 수 있고, t_max 같은 주입 스케줄도 조정 대상이다.