PRX Part 3: 24시간 만에 텍스트-이미지 모델 학습하기
도입부
최근 몇 년간 텍스트-이미지 생성 모델은 놀라운 발전을 이루었습니다. 과거에는 수백만 달러의 비용과 수주에서 수개월의 훈련 시간이 필요했지만, 이제는 훨씬 더 적은 비용과 시간으로 경쟁력 있는 모델을 만들 수 있게 되었습니다. 이번 PRX Part 3에서는 24시간이라는 제한된 시간과 예산($1,500) 내에서 텍스트-이미지 모델을 학습시키는 과정을 소개합니다. 이는 단순히 실험적인 시도가 아니라, 향후 대규모 모델 훈련의 기반이 될 수 있는 중요한 연구입니다.
이전 두 개의 게시글(Part 1 및 Part 2)에서는 다양한 아키텍처 및 훈련 트릭을 개별적으로 평가하고, 어떤 요소가 성능 향상에 기여하는지 분석했습니다. 이번에는 가장 유망한 요소들을 결합하여 성능을 극대화하는 데 초점을 맞추겠습니다. 공개된 코드(Github 링크)를 통해 다른 연구자들이 이 과정을 재현하고, 수정하고, 확장할 수 있도록 지원할 것입니다. 특히, 텍스트-이미지 모델 훈련의 효율성을 높이는 데 기여하는 기술들을 심층적으로 살펴보겠습니다.
본문
X-prediction 및 픽셀 공간 훈련
이번 훈련에는 Back to Basics에서 제안한 x-prediction 방식을 사용했습니다. 이 방식은 VAE 없이도 픽셀 공간에서 직접 훈련할 수 있도록 지원하며, 모델 구조를 단순화하는 데 기여합니다. 패치 크기는 32, 초기 토큰 투영 레이어의 병목 현상은 256차원으로 설정했습니다. 이는 시퀀스 길이를 제어하여 고해상도에서도 텍스트-이미지 훈련이 계산적으로 가능하다는 것을 의미합니다.
예를 들어, 512px 해상도에서는 시퀀스 길이가 (512/32)^2 = 256이 되고, 1024px 해상도에서는 (1024/32)^2 = 1024이 됩니다. 기존의 256px → 512px → 1024px 스케줄 대신, 512px에서 직접 시작하여 1024px로 미세 조정하는 방식을 채택했습니다. 이는 텍스트-이미지 모델 훈련의 효율성을 높이는 중요한 요소입니다.
지각적 손실 (Perceptual Losses)
픽셀 공간에서 x0를 직접 예측하는 또 다른 장점은 기존의 컴퓨터 비전 도구를 재사용할 수 있다는 것입니다. 래턴트 공간에서 모델의 출력을 디코딩하거나, 인간의 인지 능력과 일치하지 않을 수 있는 학습된 래턴트 공간에서 손실을 정의해야 하는 불편함 없이, 직접 픽셀을 예측할 수 있습니다. 이를 통해 기존의 지각적 손실을 그대로 적용할 수 있습니다.
PixelGen 논문에서 제안된 추가적인 지각적 목표를 사용했습니다. LPIPS 및 DINO 기반의 지각적 손실을 사용했으며, 이는 훈련 속도와 최종 이미지 품질을 향상시키는 데 기여합니다. 특히, 전체 풀 이미지를 사용하고, 모든 노이즈 레벨에서 지각적 손실을 적용하는 방식으로 개선했습니다. LPIPS 손실에는 0.1, DINO 지각적 손실에는 0.01의 가중치를 사용했습니다. 이러한 손실은 트랜스포머 순방향 패스와 비교적 가벼우며, 품질 향상에 도움을 줍니다. 텍스트-이미지 생성 모델의 성능을 극대화하는 데 중요한 역할을 합니다.
토큰 라우팅 (Token Routing) with TREAD
TREAD는 각 단계를 저렴하게 만드는 토큰 라우팅 기술을 사용했습니다. 이는 일부 토큰을 선택적으로 건너뛰고, 나중에 다시 주입하여 전체 정보를 유지하는 방식입니다. SPRINT보다 단순성 때문에 TREAD를 선택했습니다. TREAD는 50%의 토큰을 2번째 블록에서 마지막 블록까지 라우팅합니다.
라우팅된 모델은 CFG(Classifier-Free Guidance)에서 성능이 저하될 수 있으며, 특히 훈련이 부족할 때 문제가 발생할 수 있습니다. 따라서 Guiding Token-Sparse Diffusion Models에서 영감을 얻은 간단한 자체 가이드 스킴을 구현하여 조건 예측을 통해 모델을 안내합니다. 텍스트-이미지 모델의 안정성을 확보하는 데 중요한 기술입니다.
표현 정렬 (Representation Alignment) with REPA and DINOv3
표현 정렬에는 REPA를 사용했습니다. 교사 모델로는 DINOv3를 선택했습니다. REPA 손실은 8번째 트랜스포머 블록에서 한 번 적용하며, 손실 가중치는 0.5입니다. TREAD 라우팅과 REPA를 함께 사용하기 때문에, 정렬 손실은 라우팅되지 않은 토큰에 대해서만 계산합니다. 이는 REPA 신호의 일관성을 유지하고, 계산 경로를 건너뛴 토큰에 대한 기능 비교를 방지합니다. 텍스트-이미지 모델의 정확도를 높이는 데 기여합니다.
깊이 있는 분석
이번 24시간 훈련 결과는 텍스트-이미지 모델 개발의 현황을 잘 보여줍니다. 제한된 시간과 예산 내에서도 상당한 수준의 결과를 얻을 수 있었으며, 이는 기술 발전의 결과입니다. 아직 몇 가지 개선해야 할 부분이 있지만, 이는 훈련 데이터의 부족이나 모델 구조의 결함보다는 훈련 부족에 따른 현상으로 보입니다. 더 많은 데이터와 컴퓨팅 자원을 확보하면 이러한 문제를 해결할 수 있을 것입니다.
이번 훈련 과정을 통해 얻은 경험과 코드는 향후 대규모 텍스트-이미지 모델 개발에 활용될 것입니다. 또한, Github 저장소를 공개하여 커뮤니티의 참여를 유도하고, 연구 개발을 가속화하고자 합니다. 앞으로도 이 분야에 대한 지속적인 연구와 노력을 통해 더욱 발전된 텍스트-이미지 생성 모델을 선보일 것입니다.
결론
이번 PRX Part 3은 텍스트-이미지 모델 훈련의 가능성을 보여주는 중요한 시점입니다. 지속적인 연구와 개발을 통해 더욱 발전된 모델을 만들고, 사용자들에게 더 나은 경험을 제공할 수 있도록 노력하겠습니다.
심층 분석 및 시사점
Array
