24시간 만에 텍스트-이미지 모델 학습시키기: Photoroom의 PRX 프로젝트 3부

이 요약은 AI가 원문을 분석해 생성했습니다. 정확한 내용은 원문 기준으로 확인하세요.

핵심 요약

대규모 확산 모델(Diffusion Model) 학습에는 막대한 비용과 시간이 소요되지만 효율적인 아키텍처와 최적화 기법을 결합하면 이를 획기적으로 줄일 수 있습니다. Photoroom 팀은 32대의 H200 GPU를 사용하여 약 1,500달러의 예산으로 24시간 만에 경쟁력 있는 텍스트-이미지 모델을 학습시키는 스피드런 실험을 진행했습니다. 이 글은 VAE 없는 픽셀 공간 학습, TREAD 토큰 라우팅, REPA 표현 정렬, Muon 옵티마이저 등 성능을 극대화한 기술적 레시피를 상세히 공개합니다. 결과적으로 짧은 학습 시간에도 불구하고 강력한 프롬프트 준수 능력과 시각적 품질을 갖춘 모델을 구현했으며 관련 코드를 오픈소스로 제공합니다.

배경

확산 모델(Diffusion Models)에 대한 기본 이해, 트랜스포머(Transformer) 아키텍처 지식, PyTorch 프레임워크 사용 경험

대상 독자

이미지 생성 모델을 직접 학습하거나 최적화하려는 AI 엔지니어 및 연구자

의미 / 영향

수백만 달러가 들던 모델 학습 비용을 수천 달러 수준으로 낮출 수 있음을 증명했습니다. 이는 중소 규모 기업이나 연구실에서도 독자적인 고성능 텍스트-이미지 모델을 구축할 수 있는 길을 열어줍니다.

섹션별 상세

기존의 잠재 공간(Latent space) 방식 대신 x-예측(x-prediction) 공식을 사용하여 픽셀 공간에서 직접 학습을 진행했습니다. 패치 크기를 32로 설정하고 병목 레이어를 활용해 시퀀스 길이를 조절함으로써 1024px 고해상도에서도 연산 효율성을 유지했습니다. 이는 구조를 단순화할 뿐만 아니라 기존 컴퓨터 비전의 지각 손실(Perceptual Loss) 함수들을 직접 적용할 수 있는 장점을 제공합니다.

모델의 수렴 속도와 시각적 품질을 높이기 위해 LPIPS와 DINOv2 기반의 지각 손실을 추가했습니다. 또한 DINOv3를 스승(Teacher) 모델로 사용하는 REPA(Representation Alignment) 기법을 8번째 트랜스포머 블록에 적용하여 모델이 더 풍부한 의미적 특징을 학습하도록 유도했습니다. 이러한 보조 손실 함수들은 메인 연산 대비 오버헤드가 적으면서도 품질 향상에 크게 기여했습니다.

토큰 라우팅 기법인 TREAD를 도입하여 토큰의 50%가 중간 트랜스포머 블록을 건너뛰게 함으로써 스텝당 비용을 절감했습니다. 최적화 도구로는 Adam보다 성능이 뛰어난 Muon 옵티마이저를 2D 파라미터에 적용하고 나머지는 Adam을 사용하는 하이브리드 방식을 채택했습니다. 이를 통해 제한된 시간 내에 더 효율적인 가중치 업데이트와 빠른 성능 도달이 가능했습니다.

Flux 및 Midjourney 기반의 합성 데이터셋 약 870만 개를 사용했으며 Gemini 2.5 Flash를 통해 프롬프트를 일관성 있게 재작성(Re-captioning)했습니다. 학습은 먼저 512px 해상도에서 10만 스텝을 진행한 후 1024px에서 2만 스텝을 추가로 파인튜닝하여 세부 묘사를 강화했습니다. 최종 모델은 일부 해부학적 오류가 있을 수 있으나 전반적인 미적 완성도와 프롬프트 이행 능력이 매우 뛰어난 것으로 나타났습니다.

이미지 분석

Photo
제한된 시간과 예산 내에서 학습된 모델이 1024px 해상도에서 구현할 수 있는 시각적 디테일을 보여줍니다. 털의 질감과 눈의 표현력이 뛰어나며 제안된 학습 레시피의 실효성을 증명하는 핵심 결과물입니다.
24시간 학습 후 생성된 표범의 고해상도 이미지

실무 Takeaway

VAE를 제거한 픽셀 공간 학습은 아키텍처를 단순화하고 기존 컴퓨터 비전 손실 함수 활용도를 높입니다
TREAD와 같은 토큰 라우팅 기법은 품질 저하를 최소화하면서 학습 비용을 유의미하게 절감합니다
Muon 옵티마이저는 대규모 확산 모델 학습에서 Adam 대비 빠른 수렴 성능을 보여줍니다

언급된 리소스

GitHubPRX GitHub Repository