
PRX Part 4: 우리의 데이터 전략
요약
PRX 모델 학습을 위한 데이터 전략을 다룹니다. 공개 및 내부 데이터셋을 혼합하고 VLM을 활용해 이미지 캡션을 재작성하여 다양성 있는 사전 학습 코퍼스를 구축하는 파이프라인을 설명합니다.
핵심 포인트
- 사전 학습의 핵심은 미학적 완벽함보다 데이터의 커버리지와 다양성 확보
- VLM을 사용하여 이미지 캡션을 재작성(re-caption)하는 파이프라인 활용
- 사전 학습은 데이터의 폭을, 미세 조정은 데이터의 취향을 목표로 함
- Mosaic Data Shards(MDS)를 통한 효율적인 분산 학습 데이터 관리
한 문장으로 요약하자면: 우리는 공개 데이터셋과 내부 데이터셋을 혼합하여 학습 데이터를 수집하고, VLM (Vision-Language Model)을 사용하여 이미지의 캡션을 다시 작성(re-caption)하며, 그 결과물을 PRX 학습에 사용된 스트리밍 가능한 코퍼스(streamable corpus)로 변환합니다.
높은 수준에서 데이터 파이프라인은 다음과 같습니다:

다음에서는 이를 자세히 살펴보겠습니다.
목표는 사전 학습 (pre-training)을 위한 대규모의 다양성 있는 데이터셋을 구축하는 것이었습니다. 이 단계에서 모델은 세상이 어떻게 보이는지, 즉 시각적 개념, 사물과 장면, 사물의 구성과 조명, 그리고 이미지가 담을 수 있는 방대한 범위를 학습합니다. 이는 개별 이미지의 완벽함의 문제가 아니라, 커버리지 (coverage)와 다양성 (diversity)의 문제입니다. 광범위하고 대표성 있는 코퍼스는 개별 이미지가 평범한 스냅샷이거나 약간 압축되어 있더라도, 더 작고 아름다운 코퍼스보다 시각적 세계의 구조에 대해 모델에게 훨씬 더 많은 것을 가르쳐 줍니다. 이 단계에서 미학 (aesthetics)을 위해 과도하게 필터링하는 것은 오히려 분포를 좁히고, 나중에 회복할 수 없는 개념과 구성적 다양성을 모델로부터 앗아가는 결과를 초래하여 해가 될 수 있습니다. 생성된 결과물을 세련되게 만드는 것은 별개의 후속 과제이며, 이는 엄격하게 큐레이션된 작은 데이터셋을 통한 미세 조정 (fine-tuning) 및 선호도 정렬 (preference alignment) 단계로 남겨둡니다. 사전 학습 (pre-training)은 폭을 위한 것이고, 미세 조정 (fine-tuning)은 취향을 위한 것입니다.
우리는 공개 데이터셋과 내부 데이터셋을 혼합하여 사전 학습 (pre-training) 데이터를 구성합니다. 이 단계에서의 우선순위는 폭과 다양성, 그리고 이미 존재하는 큐레이션 (curation)을 직접 다시 하기보다는 그 위에 올라타는 것입니다. 소스 자체가 이미 품질 필터링 (quality-filtered), 중복 제거 (deduplicated), 그리고 NSFW 콘텐츠 및 개인 정보 필터링이 완료된 상태라면, 우리는 이를 대규모로 반복하는 대신 해당 작업의 결과물을 활용합니다. 소스는 다양한 형태로 들어옵니다. 어떤 것은 이미지 데이터 자체와 함께 제공되기도 하고, 어떤 것은 메타데이터와 기본 캡션 (baseline captions) 형태로 제공되어 우리가 공통된 형식으로 변환합니다. 우리는 실용적인 접근 방식을 취했습니다. 코퍼스 (corpus)를 완전히 처음부터 구축하기보다는, 기존 데이터셋과 자체 도구 (tooling)에 의존하여 빠르게 구축했습니다. 지나고 보니 이것이 구축할 수 있는 절대적으로 최고의 데이터셋은 아닐지 모르지만, 7B 모델을 사전 학습하기 위한 견고하고 가벼운 시작점이었습니다.
우리는 분산 학습 (distributed training)을 위한 데이터셋 형식으로 Mosaic Streaming과 Mosaic Data Shards (MDS)를 한동안 사용해 왔습니다. Mosaic Composer와 결합했을 때, 이는 유지보수가 매우 적게 들고 유연하며 성능이 뛰어난 분산 학습 프레임워크임을 확인했습니다. 또한, MDS 데이터셋은 쉽고 효과적으로 혼합 및 셔플링 (shuffling)할 수 있으며, S3나 GCS와 같은 객체 스토리지 (object storage)에서 직접 분산 학습을 수행할 수 있게 해줍니다.
하지만 MDS 데이터셋은 매우 경직되어 있습니다. 컬럼 (column)을 추가하거나 특정 필터를 위한 서브셋 (subset)을 만드는 것은 기본적으로 전체 데이터셋을 스캔하고 다시 써야 함을 의미합니다. 이것이 바로 우리가 이러한 종류의 피처 엔지니어링 (feature engineering)과 데이터셋 큐레이션 (curation)을 위해 Lance를 사용하는 이유입니다. Lance는 저렴한 프레디케이트 푸시다운 (predicate pushdown), 스칼라 인덱스 (scalar indexes), 그리고 벡터 검색 (vector search) 기능을 갖춘 컬럼형 데이터 형식 (columnar data format)으로, 수십억 개의 행을 가진 데이터셋을 구축하고 탐색하는 데 적합한 도구입니다.

이 두 형식은 이 포스트와 PRX 데이터 파이프라인 (data pipeline) 전반에 걸쳐 서로 보완적인 역할을 합니다. 즉, 구축을 위한 Lance, 스트리밍을 위한 MDS입니다.
이전 학습 실행(training runs)에서는 T5Gemma를 텍스트 인코더 (text encoder)로 사용하였고, 텍스트 잠재 변수 (text latents)를 미리 계산하여 MDS에 바이트 (bytes) 형태로 저장했습니다. 이번에는 텍스트 인코더를 Qwen3-VL로 전환한 후, 대신 학습 중에 실시간으로 (on the fly) 텍스트 잠재 변수를 계산하기로 결정했습니다. 학습 루프 (training loop) 내부에서 텍스트 인코더를 실행하면 처리량 (throughput) 손실이 발생하지만, 그 정도는 모델에 따라 다릅니다. 작은 디노이저 (denoiser)의 경우에는 그 차이가 상당할 수 있지만, PRX의 7B 규모에서는 텍스트 인코더의 연산량이 디노이저에 비해 무시할 수 있는 수준이며, 측정 결과 처리량 비용은 약 3~4% 정도에 불과했습니다 (30일 실행 기준 약 1일 추가).
그 대가로 우리는 두 가지를 얻습니다. 사전 계산 (pre-computation)을 건너뜀으로써 MDS 샤드 (shards)의 크기를 훨씬 작게 유지할 수 있으며, 이는 객체 스토리지 (object storage)에서 네트워크를 통해 스트리밍하는 대신 우리 SLURM 클러스터의 SSD 기반 공유 파일 시스템에 전체 사전 학습 (pre-training) 데이터셋을 저장할 수 있을 만큼 충분히 작습니다. 또한, 저장된 테라바이트 단위의 잠재 변수를 다시 작성할 필요 없이 나중에 텍스트 인코더를 자유롭게 변경할 수 있는데, 이는 우리가 Qwen3-VL로 전환할 때 수행한 방식과 정확히 일치합니다.
우리는 모든 이미지를 PNG와 같은 무손실 포맷 (lossless format) 대신 품질 92의 JPEG로 인코딩했습니다. 단순히 품질 92가 안전하다고 가정한 것이 아니라, 직접 측정했습니다. 실제 이미지는 보통 이미 여러 번 JPEG 압축이 된 상태이므로, 핵심 질문은 한 번 더 재인코딩 (re-encode)하는 것이 해가 되는지 여부입니다. 고해상도 (1–2 MP) 및 저해상도 (0.25–0.5 MP) 실제 이미지 모두에 대해 반복적인 디코드/인코딩 (decode/encode) 사이클을 수행한 결과, 품질 92에서의 첫 번째 재인코딩은 사실상 인지할 수 없는 수준이었습니다. JPEG는 빠르게 안정적인 상태로 수렴하기 때문에 각 추가 사이클은 거의 아무런 영향을 주지 않으며, 10번의 사이클 이후에도 이미지는 인지할 수 없는 범위 내에 잘 유지되는 반면, PNG는 지각적 이득 없이 3~10배 더 컸습니다. 품질 92에서 100개의 이미지를 평균 내어 원본과 비교하여 측정한 수치는 다음과 같습니다 (PSNR이 높고 LPIPS가 낮을수록 좋습니다):
| 이미지 해상도 (Image resolution) | 1× 후 PSNR (dB) ↑ | 1× 후 LPIPS ↓ | 10× 후 PSNR (dB) ↑ | 10× 후 LPIPS ↓ |
|---|---|---|---|---|
| 1–2 MP | 48.7 | 0.004 | 45.4 | 0.008 |
| 0.25–0.5 MP | 45.1 | 0.005 | 42.2 | 0.010 |
대부분의 소스 이미지가 이미 JPEG 압축된 상태로 들어오기 때문에, 이를 PNG로 무손실 (losslessly) 저장하는 것은 이득이 거의 없습니다. 따라서 우리는 모든 것을 고품질 (quality 92)의 JPEG로 변환합니다.
우리는 또한 무엇이 가장 중요한지, 즉 JPEG로 학습하는 것이 모델의 생성 결과물을 변화시키는지 확인했습니다. 어차피 우리의 코퍼스 (corpus)는 대부분 JPEG이므로, PNG가 제공할 수 있는 유일한 장점은 새로운 아티팩트 (artifacts)를 도입하지 않는 것입니다. 우리는 고해상도 소스에서 의도적으로 비교를 수행했는데, 이는 원래 JPEG 이미지라 하더라도 고해상도에서는 아티팩트가 거의 없을 것이라는 기대 때문이었습니다. 우리는 동일한 이미지에 대해 1024px에서 두 개의 동일한 PRX 모델을 학습시켰으며, 하나는 PNG로 저장된 것을 사용했고 다른 하나는 quality 92의 JPEG로 저장된 것을 사용했습니다. 두 모델 모두 우리의 지표상으로 동일하게 잘 학습되었으며, 양자화 테이블 (quantization tables)을 매칭하여 JPEG 품질을 추정하는 알려진 기술을 사용하여 출력물이 얼마나 JPEG 압축된 것처럼 보이는지 추정했을 때를 포함하여, 생성된 결과물은 사실상 구분이 불가능했습니다.
| 학습 데이터 (Model trained on) | 탐지율 (Detection rate) | 평균 추정 JPEG 품질 (Mean est. JPEG quality) | 중앙값 추정 JPEG 품질 (Median est. JPEG quality) |
|---|---|---|---|
| JPEG (quality 92) | 12.0% | 39.6 | 34.0 |
| PNG | 10.8% | 42.1 | 45.0 |
탐지율(detection rate)은 어떠한 양자화 구조(quantization structure)라도 발견된 생성물의 비율을 의미합니다. 평균 및 중앙값으로 추정된 JPEG 품질은 플래그(flagged)가 지정된 이미지들만을 대상으로 산출되었습니다. 두 모델은 사실상 구분이 불가능합니다. 두 모델 모두 약 10번의 생성 중 단 한 번 정도만 감지 가능한 양자화 구조를 보였으며, 그 차이가 매우 작기 때문에 학습 이미지 형식이 출력 품질에 미치는 영향은 미미하다고 확신합니다. 따라서 고품질 JPEG 저장은 소스 데이터가 이미 보유하고 있는 것 이상의 측정 가능한 이점을 제공하지 않습니다. 이는 충분히 훌륭한 수준의 대규모 텍스트-투-이미지 (text-to-image) 학습에는 적합합니다. 다만, 우리의 커스텀 AI Shadows 모델과 같이 아티팩트(artifact)에 민감한 다른 Photoroom 모델 학습에 사용되는 데이터의 경우, 우리는 여전히 PNG를 사용합니다.
수억 개 또는 수십억 개의 행(row)을 가진 Parquet 테이블은 대화형으로 탐색할 수 없습니다. 따라서 우리는 데이터를 인덱싱(indexing), 쿼리(querying), 브라우징(browsing)할 수 있는 Lance에 저장합니다. 우리는 Ray Data를 사용하여 인제스트(ingest)를 실행하며, 소스 테이블을 병렬로 읽어 들여 클러스터 전반에 걸쳐 Lance 테이블의 많은 프래그먼트(fragment)들을 작성합니다.
공유할 만한 교훈 중 하나는 프래그먼트화(fragmentation)에 관한 것입니다. Lance 테이블은 여러 프래그먼트로 분할되며, 몇몇 작업은 행의 수가 아닌 전체 프래그먼트 수에 따라 규모가 결정됩니다. 예를 들어, 스캔(scan) 작업은 모든 프래그먼트의 파일을 열어야 하며, 일부 메타데이터 작업은 O(프래그먼트 수)의 복잡도를 가집니다. 따라서 테이블이 너무 많은 작은 프래그먼트로 쪼개져 있으면 테이블의 크기와 상관없이 쿼리 속도가 느려집니다. Lance의 성능 가이드는 이 개수를 낮게 유지하고(10억 개의 행을 가진 테이블이라도 경험칙상 100개 정도의 프래그먼트 수준을 권장함), 작은 프래그먼트들을 더 큰 것들로 병합하기 위해 주기적으로 컴팩션(compaction)을 실행하는 것입니다.
우리는 이를 느린 방식으로 배웠습니다. 우리의 첫 번째 인제스트(ingest)는 프래그먼트(fragment)당 단 100,000개의 행(row)만을 대상으로 했으며, 이로 인해 수천 개의 아주 작은 프래그먼트들이 남게 되었고, 간단한 필터링과 전체 텍스트 쿼리(full-text query)조차 매우 느려졌습니다. 프래그먼트당 약 100만 행까지 컴팩션(compaction)을 수행한 결과(수억 개의 행을 가진 우리 테이블의 경우 약 1,000개의 프래그먼트 수준), 쿼리는 빨라졌고 더 이상 이 문제를 걱정하지 않게 되었습니다. 적절한 목표치는 행의 너비(row width), 쿼리 패턴, 그리고 데이터가 기록되는 방식에 따라 달라지며, 지나고 보니 더 적고 더 큰 프래그먼트가 더 나은 기본값이었을 수도 있습니다. Lance는 분산 컴팩션(distributed compaction, 예를 들어 Lance-Ray 통합을 통해)을 지원하므로 사후에 조정하기는 쉽지만, 애초에 과도하게 프래그먼트화(over-fragment)하지 않는 것이 좋습니다.

우리는 데이터셋에 포함되어 있는 캡션(caption)에 의존하기보다 모든 이미지의 캡션을 직접 다시 작성하기로 결정했습니다. 그 이유는 일관성 때문입니다. 다양한 출처가 섞여 있으면 캡션의 길이, 스타일, 품질이 크게 달라지기 때문에, 전체 코퍼스(corpus)에 걸쳐 단일하고 균일한 표준을 유지하고 싶었습니다(캡셔너(captioner)에 대한 자세한 내용은 섹션 3에서 다룹니다).
그럼에도 불구하고 데이터셋이 이미 보유하고 있는 메타데이터(metadata)는 다른 용도로 여전히 가치가 있습니다. 기존의 캡션과 때때로 그와 함께 제공되는 벡터 임베딩(vector embeddings, 예: CLIP 임베딩)은 Lance에서 데이터셋을 즉시 탐색 가능하게 만드는 요소입니다. 캡션에 대한 전체 텍스트 검색(full-text search)과 임베딩 컬럼(embedding column)에 대한 최근접 이웃 검색(nearest-neighbor search)을 통해, 자체적인 캡셔닝을 실행하기 훨씬 전부터 데이터를 탐색하고 품질 및 필요한 필터링을 빠르게 판단할 수 있습니다. 따라서 우리는 두 가지를 모두 유지하며, 별도의 테이블에 있는 경우 이를 조인(join)하여 사용하고, 다음에 설명할 브라우징(browsing) 과정에서 이를 활용합니다.
데이터를 쿼리(query)할 수 있게 됨에 따라, 우리는 이를 빠르게 프로파일링(profiling)할 수 있었습니다. 예를 들어, 해상도 분포를 통해 컷오프(cutoff)를 어디에 설정할지 결정할 수 있었습니다. 우리의 가장 작은 학습 버킷(training bucket)은 512² 픽셀이며, 버킷으로의 업스케일링(upscaling)은 4/3(약 33%)로 제한하므로, 유효한 컷오프는 384² ≈ 147k 픽셀과 [0.5, 2.0] 범위의 종횡비(aspect ratio)가 됩니다. 이보다 낮은 모든 데이터는 폐기합니다.
Lance 테이블은 텍스트 컬럼에 대한 전체 텍스트 검색(full-text search)뿐만 아니라 벡터 임베딩(vector embedding) 컬럼에 대한 최근접 이웃 유사도 검색(nearest-neighbor similarity search)을 기본적으로 지원하며, 두 기능 모두 인덱스(index)를 구축하여 가속화할 수 있습니다. 우리는 캡션(caption)과 임베딩 컬럼에 정확히 이 방식을 적용한 후, 데이터셋을 브라우징(browsing)할 수 있는 작은 UI를 구축했습니다. 이를 통해 캡션에 대한 실시간 텍스트 검색과 벡터 검색을 통한 시각적으로 유사한 이미지 클러스터(cluster) 탐색이 가능해졌습니다.
이러한 방식으로 데이터를 브라우징함으로써, 우리는 단순히 데이터를 살펴보는 것만으로도 데이터의 다양성과 품질을 정성적으로 평가하고 문제를 조기에 발견할 수 있었습니다. 우리는 정보가 없는 기본 캡션, 놀라울 정도로 많은 비사진적 콘텐츠(스크린샷, 슬라이드, 문서, 인포그래픽), 그리고 일부 거의 중복된 이미지들을 발견했습니다. 이러한 관찰 결과는 파이프라인의 나머지 부분을 직접적으로 형성했습니다. 즉, 모든 것에 대해 재캡션(re-caption)을 수행하고(더 길고 정확한 캡션은 불완전한 이미지조차 잘 설명하며, 우리가 학습하는 정확한 이미지와 캡션이 일치함을 보장합니다), 나중에 가벼운 필터링(filtering) 및 중복 제거(deduplication) 단계를 추가했습니다. 일반적으로, 우리는 이러한 종류의 탐색 도구가 데이터셋을 확정하기 전에 데이터의 느낌을 파악하는 데 매우 귀중하다고 믿습니다.
우리는 길고 상세한 캡션이 출력 품질에 매우 강력한 레버(lever)가 된다는 것을 발견했습니다. 초기 벤치마크에서 우리는 동일한 이미지로 작은 확산 모델(diffusion model)을 두 번 학습시켰습니다. 한 번은 Qwen2.5-VL-7B로 생성한 캡션을 사용했고, 다른 한 번은 훨씬 짧은 캡션(LLaVA-1.5-LLaMA3-8B로 생성)을 사용했으며, 학습 과정 전반에 걸쳐 두 모델을 평가했습니다. Qwen2.5-VL-7B로 생성된 긴 캡션이 모든 체크포인트에서 더 낮은 FID, CMMD, DINO-MMD를 기록하며 전반적으로 승리했습니다(세 지표 모두 낮을수록 좋습니다).

Qwen2.5-VL-7B 캡션(빨간색)으로 학습된 소규모 확산 모델 (diffusion model) 대 베이스라인인 LLaVA-1.5-LLaMA3-8B 캡션(파란색)을 비교한 결과로, 100k 스텝까지 매 10k 스텝마다 점수를 측정했습니다. 세 가지 지표 모두 낮을수록 좋습니다.
최종 지표 (약 100k 스텝):
AI 자동 생성 콘텐츠
본 콘텐츠는 HuggingFace Blog의 원문을 AI가 자동으로 요약·번역·분석한 것입니다. 원 저작권은 원저작자에게 있으며, 정확한 내용은 반드시 원문을 확인해 주세요.
원문 바로가기