ByteDance Seed 연구팀이 잠재 공간 VAE와 생성 모델을 함께 학습하는 방법 GenFirst의 가중치와 코드를 Hugging Face와 GitHub에 공개했다. 가중치는 CC BY-NC 4.0 라이선스라 비상업 연구에만 쓸 수 있고, 코드는 MIT 라이선스다. 논문은 arXiv:2608.29335로 올라 있다.
이번에 풀린 것은 확산 모델(diffusion) 방식의 EiT 구현이다. 논문에 실린 자기회귀 방식의 EAR 구현은 공개 대상이 아니다. 공개물은 종단간(end-to-end) 방식으로 학습한 VAE와, 그 VAE를 고정한 채 학습한 생성 모델 두 종류다. 후자는 클래스 조건 이미지 생성용 SiT와 텍스트-이미지 생성용 MMDiT다. 저자는 Guangting Zheng, Yiyuan Zhang, Tao Yang, Yunpeng Chen, Rui Zhu, Jiajun Deng, Yanyong Zhang이다.
학습 방식
VAE와 생성 모델을 함께 학습하면 잠재 공간이 붕괴해 학습이 불안정해지기 쉽다. 연구팀은 GenFirst가 이 붕괴 없이 두 모델을 함께 학습한다고 설명했다. 방법은 두 단계 일정이다. 1단계에서는 생성 모델(prior)이 잠재 공간의 구조를 잡게 하고, 2단계에서는 prior 가중치와 엔트로피 가중치를 모두 10분의 1로 낮춰 디코더가 세부 묘사를 되찾게 한다. 그 결과 확산 모델이 잘 맞추면서도 오토인코더로 쓸 만한 잠재 공간이 만들어진다고 한다.
종단간 단계는 전역 배치 256으로 높은 가중치에서 40만 스텝, 낮은 가중치에서 20만 스텝을 학습한다. 이후 VAE를 고정하고 잠재값을 캐시한 뒤, 실제로 평가하는 prior를 처음부터 256 해상도 400만 스텝, 512 해상도 200만 스텝 학습한다. 두 단계 사이에 선택적으로 디코더만 1만 스텝 미세 조정할 수 있고, 이 과정을 거치면 가이드 없는 gFID가 약 0.2 개선된다. 학습 데이터는 공개 데이터셋이며 개별 목록은 밝히지 않았다.
공개된 모델과 성능
SiT-XL과 MMDiT-XL은 28블록, 너비 1152, 16 어텐션 헤드로 약 6억 7,500만 개 매개변수다. 텍스트 인코더는 수정하지 않은 CLIP-L이고, 가중치는 fp32로 공개됐으며 샘플링은 BF16으로 시험했다. 클래스 조건 모델은 0~999 정수 라벨 하나만 입력받고 텍스트는 받지 않는다.
| 학습 에포크 | 해상도 | gFID(가이드 없음) | gFID(가이드 적용) | IS(가이드 적용) |
|---|---|---|---|---|
| 80 | 256x256 | 2.79 | 1.38 | 271.3 |
| 480 | 256x256 | 1.60 | 0.988 | 270.7 |
| 800 | 256x256 | 1.45 | 0.969 | 283.4 |
| 80 | 512x512 | 해당 없음 | 1.63 | 278.9 |
| 400 | 512x512 | 해당 없음 | 1.23 | 286.1 |
256 해상도 모델은 f16d32 VAE를 쓰며, 512 모델은 이 VAE를 그대로 고정하고 prior만 512에서 학습했다. 512 모델은 가이드를 적용한 결과만 보고했고, 가이드 없는 동작은 조정하지 않았다. 평가는 샘플 5만 장을 torch-fidelity로 채점했고, 샘플링은 250스텝 SDE다. 가이드 행은 같은 학습의 초기 체크포인트에서 멀어지도록 하는 autoguidance에 CFG를 더했으며, 512 행만 CFG만 썼다.
텍스트-이미지 모델은 512x512에서 MMDiT-XL/2와 고정 CLIP-L, 종단간 학습한 FLUX VAE를 쓴다. 프롬프트 재작성 없이 잰 전체 점수는 변형 A가 GenEval 0.8987, DPG-Bench 82.60이고 변형 B가 0.8920, 82.06이다. 두 변형은 학습에 쓴 VAE만 다르다.
이해와 생성이 잠재 공간을 공유하는 모델도 있다. Qwen3-VL ViT 오토인코더를 종단간으로 다듬은 것으로, SigLIP 지도를 쓴 버전은 rFID 0.39, PSNR 27.64, GenEval 0.90, DPG-Bench 83.76, 선형 프로브 Top-1 79.25%다. VLM-NLL 지도 버전은 rFID 0.46, PSNR 27.19, GenEval 0.89, DPG-Bench 83.49, 선형 프로브 81.69%다.
텍스트와 이미지를 한 모델에서 다루는 JointMMDiT-XL/2도 있다. 텍스트 VAE(Cola)와 이미지 VAE(FLUX.1)를 함께 다듬었고, 텍스트-이미지 생성은 GenEval 86.57, DPG-Bench 72.68을 기록했다. COCO 캡션 생성은 CIDEr 33.28, BLEU-4 7.76, METEOR 17.74, ROUGE-L 34.16, SPICE 11.81, CLIPScore 57.97이다. 생성에는 3단계 SFT 체크포인트를, 캡션에는 SFT 이전 체크포인트를 쓴다.
| VAE | 압축/채널 | 초기화 출처 | rFID | PSNR | SSIM | LPIPS |
|---|---|---|---|---|---|---|
| f8d4 | 8x, 4ch | SD-VAE | 0.57 | 25.13 | 0.71 | 0.14 |
| f16d32 (디코더 미세 조정, 256) | 16x, 32ch | VA-VAE | 0.36 | 25.82 | 0.74 | 0.12 |
| f8d16 변형 A | 8x, 16ch | FLUX.1 VAE | 0.23 | 29.43 | 0.85 | 0.07 |
| f8d16 변형 B | 8x, 16ch | FLUX.1 VAE | 0.21 | 29.37 | 0.85 | 0.07 |
| f8d32 (512) | 8x, 32ch | FLUX.2 VAE | 0.041 | 36.46 | 0.953 | 0.026 |
표의 수치는 사후 평균(posterior mean)으로 복원한 품질이며, 각 VAE를 실제로 쓰는 해상도에서 쟀다. f16d32 디코더 미세 조정 모델은 256, FLUX.2 VAE 기반 f8d32는 512 해상도 기준이다.
rFID는 학습셋이 아니라 5만 장 검증 이미지를 기준으로 쟀기 때문에 다른 방식으로 낸 수치와 직접 비교할 수 없다고 연구팀은 밝혔다. FID는 기준 통계와 크기 조정 방식에 민감해서, 다른 연구와의 비교는 같은 평가 절차에서만 의미가 있다는 설명도 붙였다.
쓰는 법과 제한
체크포인트는 raw.pt 형식이고 diffusers 연동이나 from_pretrained 진입점이 없다. 저장소를 내려받아 requirements를 설치한 뒤 Prior_only 폴더의 generate.py로 샘플을 만든다. SiT-XL/2는 f8d4 VAE에, SiT-XL/1은 f16d32 VAE에 맞춰 써야 한다. 텍스트-이미지 샘플링 설정과 전체 평가 코드는 저장소 docs 폴더에 있다.
VAE는 반드시 짝이 되는 잠재 통계 파일과 함께 써야 한다. prior가 그 통계로 정규화한 잠재값으로 학습됐기 때문에, 엉뚱한 통계 파일을 물리면 품질이 떨어지는 정도가 아니라 노이즈가 나온다. 오류 없이 출력이 나오므로 짝을 직접 확인해야 한다. FLUX.2 기반 f8d32 VAE는 safetensors 형식과 채널별 scaling_factor, shift_factor 값이 함께 제공된다.
.pt 체크포인트는 pickle 형식이라 불러오면 임의의 Python 코드가 실행될 수 있다. 인자 네임스페이스가 함께 저장돼 있어 torch 2.6 이상에서도 weights_only=False로 불러와야 하므로, 공식 저장소에서 받은 파일만 써야 한다고 안내했다.
라이선스는 겹겹이다. 가중치는 CC BY-NC 4.0이고 코드는 MIT지만, 저장소 전체로는 상업 이용이 어렵다. 데이터 전처리가 NVIDIA EDM2(CC BY-NC-SA 4.0)에서 나왔고, 일부 인코더와 샘플러는 Meta 프로젝트(CC BY-NC 4.0)에서 이어졌다. 구성요소별 내역은 저장소의 NOTICE에 있다. 가중치는 stabilityai/sd-vae-ft-ema, LightningDiT의 VA-VAE, black-forest-labs의 FLUX.1 VAE, OpenAI CLIP 같은 상위 모델에서 파생돼 각각의 약관이 계속 적용된다. 학습 데이터셋의 약관도 별도로 적용된다.
허용 용도는 논문 재현, 종단간 VAE와 prior 학습 연구, 잠재 공간 분석, 비상업 연구의 기준선이나 초기값이다. 상업 이용, 실존 인물 생성, 성적·폭력적·혐오 이미지 생성, 실제 사진이나 사실 기록처럼 제시하는 용도, 안전이 중요하거나 의사결정에 쓰는 파이프라인은 범위 밖이다. 안전 필터, 워터마크, 콘텐츠 분류기는 들어 있지 않아서, 출력이 다른 사람에게 닿는 환경에 배포하려면 이용자가 직접 통제 장치를 붙여야 한다.
한계도 명시돼 있다. 클래스 조건 모델은 고정된 1,000개 라벨만 다루고 텍스트를 따르지 못하며, 큐레이션된 이미지 모음에 있는 지역·문화 쏠림이 출력에 반영된다. 텍스트-이미지 모델은 학습 데이터와 CLIP-L의 편향을 물려받아 인구집단·직업·문화별 표현이 고르지 않을 수 있고, 길거나 구성이 복잡하거나 개수를 세야 하는 프롬프트에서 지시 이행이 떨어진다. 글자, 손, 미세한 구조는 같은 규모의 비슷한 모델처럼 여전히 불안정하다.
연구팀은 GenEval과 DPG-Bench 점수가 해당 벤치마크의 프롬프트 정합성만 재며 안전성이나 공정성을 보증하지 않는다고 밝혔다. FID 같은 분포 수준 지표는 개별 샘플의 품질이나 적절성을 말해 주지 않는다는 설명도 덧붙였다. 문의는 GitHub Issues나 저자 이메일로, 보안 취약점 신고는 코드 저장소의 SECURITY.md에 안내된 ByteDance 보안팀 경로로 하라고 안내했다.
출처
이 기사는 아래 자료를 바탕으로 정리했습니다. 제목을 누르면 원문이 새 창으로 열립니다.
- ByteDance-Seed/GenFirst 새 저장소/공개 관측(Hugging Face) huggingface.co