PRX 데이터 전략: VLM 재캡셔닝 및 Mosaic Streaming을 통한 사전 학습 확장

Photoroom은 7B 파라미터 모델인 PRX를 학습하는 데 사용된 데이터 전략에 대한 상세한 분석 내용을 공개했습니다. 핵심 접근 방식은 공개 및 내부 소스로부터 다양한 코퍼스를 수집하고, 고정밀 설명을 보장하기 위해 Vision Language Model (VLM)을 사용하여 모든 이미지를 재캡셔닝하며, 데이터셋 탐색과 학습 처리량을 균형 있게 유지하기 위해 이중 형식 저장 시스템을 활용하는 것입니다.

사전 학습 원칙: 완벽함보다는 폭넓게

사전 학습 단계에서 Photoroom은 개별 이미지의 미학보다는 커버리지와 다양성을 우선시했습니다. 목표는 시각적 세계의 광범위한 분포에 걸쳐 모델에게 시각적 개념, 구도 및 조명을 가르치는 것이었습니다.

• 다양성 vs 취향: 팀은 사전 학습 중에 미학을 위해 과도하게 필터링하는 것이 분포를 좁히고 모델에게 필수적인 구도적 다양성을 잃게 만든다고 주장합니다. 생성물의 다듬기 작업은 더 작고 큐레이션된 세트에서의 미세 조정(fine-tuning) 및 선호도 정렬(preference alignment) 단계로 미뤄집니다.
• 실용적인 소싱: 데이터는 공개 및 내부 데이터셋의 혼합으로 수집되었으며, 프로세스를 가속화하기 위해 가능한 경우 기존의 큐레이션(NSFW 및 PII 필터링)을 활용했습니다.
• 캡셔닝 철학: 길고 정확한 캡션이 매우 중요합니다. 로고, 텍스트 또는 아티팩트를 포함하여 이미지의 모든 것을 설명함으로써, 이러한 요소들은 무조건적인 노이즈가 아니라 조건화되고 제어 가능한 속성이 됩니다.

데이터 인프라: Lance 및 Mosaic Data Shards (MDS)

Photoroom은 데이터셋의 라이프사이클을 처리하기 위해 두 가지 별도의 데이터 형식을 사용합니다: 큐레이션을 위한 Lance와 스트리밍을 위한 MDS입니다.

탐색을 위한 Lance

Lance는 수십억 개의 행을 가진 데이터셋을 구축하고 탐색하기 위한 컬럼형 데이터 형식으로 사용됩니다. 이는 저렴한 predicate pushdown, scalar indexes 및 vector search를 지원하여 팀이 다음과 같은 작업을 수행할 수 있게 합니다:

• 데이터 프로파일링: 해상도 분포를 분석하여 컷오프를 설정합니다 (예: 384² 픽셀 미만의 이미지 폐기).
• 대화형 브라우징: 커스텀 UI를 통해 전체 텍스트 검색 및 최근접 이웃 유사도 검색을 사용하여 다양성을 질적으로 평가하고 비사진적 콘텐츠와 같은 문제를 식별합니다.
• 프래그먼트 관리: 팀은 Lance 테이블을 너무 잘게 나누면(예: 프래그먼트당 100k 행) 쿼리 속도가 느려진다는 점을 확인했습니다. 프래그먼트당 약 100만 행으로 압축하여 성능을 크게 향상시켰습니다.

학습을 위한 MDS

Mosaic Data Shards (MDS)는 분산 학습에 사용되며, 결정론적 재개 가능 셔플링(deterministic resumable shuffling), 탄력적 에포크 중간 체크포인팅(elastic mid-epoch checkpointing) 및 여러 스트림의 가중치 혼합을 제공합니다. 파이프라인을 최적화하기 위해:

• On-the-fly Latents: 텍스트 잠재 변수(text latents)를 미리 계산하는 대신, Photoroom은 학습 중에 Qwen3-VL 인코더를 사용하여 이를 계산합니다. 이를 통해 MDS 샤드 크기를 줄이고 테라바이트 단위의 데이터를 다시 쓰지 않고도 인코더를 교체할 수 있었으며, 7B 규모에서 처리량 손실은 3–4%에 불과했습니다.
• JPEG 인코딩: 모든 이미지는 품질 92의 JPEG로 저장됩니다. 실증 테스트 결과, 이는 PNG와 시각적으로 구별할 수 없으며 생성된 이미지의 품질에 부정적인 영향을 미치지 않으면서 저장 요구 사항을 3–10배 줄이는 것으로 나타났습니다.

VLM 재캡셔닝 전략

Photoroom은 길고 상세한 캡션이 샘플 품질을 크게 향상시킨다는 것을 발견했습니다. 벤치마크에서 Qwen2.5-VL-7B 캡션으로 학습된 작은 확산 모델은 짧은 LLaVA-1.5-LLaMA3-8B 캡션으로 학습된 모델보다 FID, CMMD 및 DINO-MMD 메트릭 전반에서 더 나은 성능을 보였습니다.

캡셔너 선택

여러 후보를 벤치마킹한 후, Photoroom은 Qwen3-VL-8B를 기본 캡셔너로 선택했습니다. 이 선택은 품질과 처리량의 균형을 바탕으로 이루어졌습니다:

• 성능: Qwen3.5-9B가 메트릭에서 가장 좋은 성능을 보였지만, Qwen3-VL-8B는 그에 근접하면서도 훨씬 빨랐습니다 (H200 GPU당 20 img/s).
• 프롬프팅: 팀은 추측이나 해석 없이 정확한 시각적 근거(visual grounding)에 집중하며, 100–200 단어의 단일 흐르는 단락을 요구하는 밀도 높은 시스템 프롬프트를 사용했습니다.

데이터셋 정제 및 중복 제거

VLM 캡션이 생성되면, Photoroom은 전체 데이터셋을 다시 쓰지 않고 코퍼스를 정제하기 위해 가벼운 필터링과 중복 제거를 적용했습니다.

캡션 기반 필터링

비싼 픽셀 수준 분류기를 사용하는 대신, Photoroom은 Qwen3-8B를 텍스트 전용 모드로 사용하여 캡션을 "visual", "text", "nsfw"로 분류했습니다. 이 접근 방식은 계산 효율적이며(GPU당 200 captions/s) 스크린샷, 문서 및 인포그래픽을 효과적으로 제거합니다.

지각 해시(Perceptual Hash) 중복 제거

유사 중복 항목을 제거하기 위해 팀은 DCT 기반의 지각 해시를 구현했습니다. 두 이미지의 Hamming 거리가 0이면 중복으로 간주됩니다. 서로 다른 해상도에 중복이 존재하는 경우, 가장 높은 해상도의 복사본이 유지됩니다.

스킵-리스트(Skip-List) 메커니즘

모든 필터링 또는 중복 제거 패스마다 거대한 MDS 코퍼스를 다시 쓰는 것을 피하기 위해, Photoroom은 데이터 로더에 스킵-리스트 기능을 구현했습니다. 샤드당 사이드카 파일이 건너뛸 인덱스 목록을 보유하여, 로드 시점에 동적으로 데이터를 제외(예: 사용자 옵트아웃 또는 품질 문제)할 수 있습니다.

종횡비 버케팅(Aspect-Ratio Bucketing)

이미지당 일정한 계산량을 유지하고 불필요한 패딩이나 크롭을 피하기 위해, Photoroom은 종횡비 버케팅을 사용합니다:

1. 해상도 티어: 이미지에는 픽셀 수에 따라 티어(512px, 1024px, 2048px, 4096px)가 할당됩니다.
2. 종횡비 버킷: 각 티어 내에서 이미지는 일정한 패치 토큰 예산을 유지하는 13개의 패치 정렬 모양(AR 0.5 ~ 2.0) 중 하나로 스냅됩니다.
3. 처리: 이미지는 Lanczos 필터를 사용하여 크기가 조정되며, 별도의 스트리밍을 위해 해상도와 종횡비로 키가 지정된 디렉토리 트리 내에 저장됩니다.