REPRESENTATION ALIGNMENTFOR GENERATION : TRAINING DIFFUSION TRANSFORMERS IS EASIER THAN YOU THINK
논문을 읽고
Jul 13, 2025
Diffusion에서 denoising process 과정의 중간 representation이 유의미한 rep이라는 연구들이 있다. 그러나 여전히 self-supervised로 vision model이 배운 것보다는 부족하긴하다.
우리는 큰 diffusion model의 생성에 있어 bottleneck이 이 representation을 만드는 능력에 있다고 생각한다.
여기서 diffusion model이 알아서 배우게 두는게 아니라 다른 고퀄리티 visual representation을 사용해서 도와줄 수 있다. 그럼 전체 generation 성능과 속도가 증가한다.
by straightforward regularization called REPA - 디노이징 과정의 noisy input hidden state와 clean image의 다른 모델에 의한 rep을 align한다..? (레이어 하나 태워서 맞춤)
본 연구에서는 DINOv2(Oquab et al., 2024)를 썼다. 제일 최신이고 좋다고 보는듯
alignment 측정은 The platonic representation hypothesis의 방법을 차용했다.
c를 보면 Diffusion model이 학습될수록 점점 alignment가 높아지긴 한다.. 대신 너무 느리고 한계가 있다.
그러나 그렇다고 off-the-shelf self-supervised vision encoder를 바로 활용하려고하면 몇가지 문제가 있다.
- DiNO는 깨끗한 이미지를 입력으로 받아서 rep을 만든다. diffusion은 noise가 추가된 이미지를 받는다.
- LDM은 애초에 latent space에서 돌기 때문에 DINO의 입력/출력(pixel image)과 다르다.
- DINO는 recon이나 generation을 위해 학습된게 아니다.
일단 관찰한 결과 diffusion의 중간 layer의 representation은
- 의미있는 표현을 배우긴 한다(특히 20th layer가 제일 높았다). 대신 DINO보다 많이 부족하다.
- MAE 보다는 낫다.
- 더 오래 학습할수록 alignment가 높아지긴한다. 그래도 많이 별로지만
따라서 핵심을 한문장으로 하자면!
REPA는 이미지 x의 pre-trained self-supervised visual representation인
y를 noise input의 diffusion transformer representation인 h로 distill한다.이렇게하면 학습에 도움이 된다는 생각을.. 할 수가 있나?
실제로는 트랜스포머 블럭이 24레이어라고 하면 앞쪽의 몇개의 블럭만 align해도 충분히 좋았다.
그럼 일단 앞에서 align 시키고 뒤에서는 생성을 위해 뭔가 디테일을 추가하는 형태로 학습하게 된다.
align은 cosine similarity로 했다. 다른 NT-Xent CE를 쓰는 것도 시도해봤지만 크게 차이 없었다.
학습은 패치간 유사도를 최대화하는 형태다.
SD VAE를 썼다. 패치 사이즈는 2. iteration이 몇인지를 알려주는데 전부 기준은 batch size 256
we observe that sufficient representation alignment can be achieved by aligning only the first few transformer blocks. This, in turn, allows the later layers of the diffusion transformers to focus on capturing high-frequency details based on the aligned representations
regularization 람다값은 0.5가 좋았다고는 함
궁금한건 diffusion transformer의 output은 image latent 자체가 아니라 vector field 잖아. 그래도 이렇게 하는게 괜찮은가?
→ 앞 8개의 레이어에만 주기때문에 괜찮다고 볼 수 있을 것 같기도 하고 아닐 것 같기도 하고. 마지막 레이어를 align시킨다고하면 오히려 이상한데
사견 : 그럼 flow matching이라고 했을 때 학습 objective자체가 일단 초기 레이어들에서는 target image의 representation을 예측하고 그 뒤로는 그걸 바탕으로 어떻게 가야하는지 vector field를 예측하는 형태가 된다? 여기서 약간 이상하다. predict하는게 target 분포가 아니라 vector field인데.. 차라리 consistency model 같은데서 뭔가 더 좋을 것 같다. representation이라는게 semantic한 정보니까 너무 그렇게 볼 필요는 없는건가
그래도 이렇게 많이 좋아지는게 신기하긴함.
Share article
Subscribe to our newsletter
