디노이징 매 스텝 예측 임베딩을 조건으로 쓰는 이미지 생성
Embedding Prediction Helps Image Generation
무엇인가
이 논문이 푸는 문제는 확산 트랜스포머의 조건 인터페이스가 지나치게 정적이라는 점이다. DiT 계열 모델은 클래스 레이블이나 텍스트 프롬프트를 한 번 임베딩하고, 그 동일한 조건을 모든 디노이징 스텝에서 재사용한다. 이미지 상태는 노이즈에서 깨끗한 샘플로 계속 변하는데 조건은 고정되어 있고, 그 조건이 현재 이미지에 대해 무엇을 의미하는지 해석하는 일은 전적으로 생성기에 맡겨진다. 저자들은 조건 자체를 현재의 노이즈 이미지로부터 매 스텝 예측할 수 있는지 묻는다.
어떻게 동작하나
핵심 부품은 NEPA(Next-Embedding Predictive Autoregression)다. NEPA는 픽셀이나 이산 토큰을 복원하는 대신 시퀀스에서 다음 연속 임베딩을 예측하도록 트랜스포머를 학습시킨다. 구조상 출력이 아직 보지 못한 임베딩의 추정치이므로 그 자체가 예측기다. 생성에서는 시퀀스가 조건 → 노이즈 이미지 → 깨끗한 이미지 순으로 흐르므로, 깨끗한 이미지의 임베딩이 바로 '다음 임베딩'이 된다. Multi-Embedding Prediction(MEP)은 이 NEPA를 확장해 다음 K개 임베딩을 한 번의 순전파로 동시에 예측한다. 이미지 생성에서는 K=N으로 두어 깨끗한 이미지의 패치 임베딩 N개를 전부 한꺼번에 예측한다. 시퀀스는 패치 위치가 아니라 생성 상태 순서로 배열되므로, 노이즈 패치 i 다음에 오는 임베딩은 깨끗한 패치 i가 된다.
무엇과 다른가
Embedding Conditioned Generation(ECG)은 이 예측을 생성기의 조건으로 쓰는 단계다. NEPA 모델을 동결한 채, 매 디노이징 스텝마다 현재 노이즈 입력을 NEPA에 넣어 예측 임베딩을 얻고, DiT 생성기는 그 임베딩만을 조건으로 받는다. 생성기에는 별도 클래스 임베딩이 없고, 타임스텝은 adaLN 변조로 들어간다. 어텐션 마스크에서 조건 토큰과 begin-of-image 토큰은 이미지 토큰을 전혀 참조하지 않는 causal prefix를 이루므로, 그 KV 캐시는 디노이징 루프 전에 한 번 계산해 매 스텝 재사용하고 이미지 토큰만 다시 인코딩한다. 손실은 InfoNCE로, 예측 임베딩을 쿼리로, 깨끗한 임베딩을 키로 두고 대각선이 양성, 같은 이미지의 다른 패치들이 음성이 된다. 타깃은 모델 자신의 패치 임베딩 레이어 출력이며 그래디언트가 타깃으로 흐르게 해야 학습이 수렴한다.
어떻게 쓰나
ImageNet-1K 256×256, SD-VAE 잠재공간(32×32×4)에서의 소규모 절제 실험이 설계 선택을 정당화한다. 손실 비교에서 InfoNCE가 FID 31.36으로 MSE 35.23, 코사인 유사도 37.82를 앞섰고, 흥미롭게도 NEPA 모델의 ImageNet 분류 정확도는 반대 순서였다. 타깃 정규화는 FID를 31.36에서 34.76으로 악화시켜 원본 타깃을 쓴다. NEPA에 타임스텝을 넣는 것은 31.36 대 31.40으로 사실상 차이가 없어 타임스텝 없는 모델을 유지한다. InfoNCE 후보 뱅크는 인스턴스 31.36, 배치 32.13, 전체 31.21이었고 통신이 필요 없는 인스턴스 뱅크를 택했다. 패치 크기는 4가 64토큰으로 31.00을 기록해 패치 2(31.36)보다 좋고 스텝당 쿼리 비용도 싸며, 패치 8은 33.20으로 너무 거칠다.
전제와 한계
조건 비교 실험이 이 논문의 핵심 주장을 가장 직접적으로 뒷받침한다. 생성기와 학습 레시피를 고정하고 조건만 바꾼 결과, 클래스 임베딩은 36.39, 클래스 토큰에 노이즈 패치 임베딩을 더한 것은 37.21이었다. 같은 NEPA-XL 아키텍처의 추가 네트워크를 붙여 파라미터와 스텝당 FLOPs를 맞춘 세 행에서는, 생성기와 end-to-end로 함께 학습한 경우 29.52, 동일 예산 240에폭으로 flow matching으로 사전학습 후 동결한 특징이 30.86, MEP로 학습해 동결한 예측 임베딩이 25.04를 기록했다. 스케일링에서는 NEPA-B/L/XL과 DiT-B/L/XL의 9개 조합 모두에서 FID가 단조 감소했고, 생성기를 B에서 XL로 키우면 모든 NEPA 모델에서 약 14점, NEPA를 B에서 XL로 키우면 모든 생성기에서 약 5점이 개선됐다. 최종 모델 NEPA-DiT-XL은 REPA 정렬 손실(DINOv2-B, 세 번째 two-stream 블록, λ=0.5)과 결합해 250스텝 SDE 샘플러로 FID 1.32, 96스텝 ODE로 1.57을 달성했으며, 총 1.39B 파라미터를 3.1×10^20 FLOPs로 학습해 800에폭에 걸쳐 1.42에 도달한 SiT-XL/2+REPA의 약 3분의 1 연산에 해당한다. 가이던스는 t∈[0.4,1] 구간에 스케일 3.6을 적용한 설정이 두 샘플러 모두에서 최선이었다.
개발자 관점에서 이 논문의 실무적 의미는 조건 인터페이스를 학습 가능한 예측 모듈로 분리할 수 있다는 것이다. 조건을 매 스텝 다시 계산하는 것이 비용만 늘리는 낭비가 아니라 FID를 실제로 끌어내린다는 수치가 있고, NEPA 스텝당 92 GFLOPs는 생성기 309 GFLOPs의 약 30% 수준이며 96 ODE 스텝 기준 이미지 한 장이 62 TFLOPs로 SiT-XL/2+REPA의 250 SDE 스텝 91 TFLOPs보다 오히려 낮다. 다만 쿼리 주기를 얼마나 자주 잡느냐가 결정적이다. 첫 스텝에서 한 번만 예측하고 재사용하면 FID가 242.63으로 붕괴하고, 8스텝마다 질의하면 1.82, 매 스텝 질의해야 1.57이 나온다. 조건 캐싱과 배치 구성, 그리고 조건 prefix를 한 번만 인코딩하는 마스크 설계를 그대로 따라 해야 이 비용 구조가 재현된다.
저자들이 명시한 전제와 한계는 분명하다. NEPA 모델은 모든 샘플링 스텝마다 두 번째 네트워크를 추가하므로, 학습과 추론 모두 네트워크 두 개를 돌린다는 점이 근본적인 오버헤드다. 그럼에도 최종 모델이 REPA의 약 3분의 1 학습 연산으로 더 나은 FID에 도달한다는 것이 이 논문의 절충이다. 또한 실험은 클래스 조건부 ImageNet 256×256 단일 설정에 한정되어 있어, 텍스트 조건부 생성이나 다른 해상도·데이터셋에서 같은 이득이 나오는지는 원문에 수치가 제시되지 않았다. 타깃 그래디언트를 끊으면 학습이 아예 수렴하지 않는다는 점, 타깃 정규화가 오히려 해롭다는 점도 재현 시 반드시 지켜야 할 조건으로 보고된다.