2만 달러로 만든 2B 이미지 생성 모델 P-JiT, VAE 없이 픽셀 공간에서 학습하고 Apache 2.0으로 공개

들어가며

이미지·영상 생성 모델을 학습하려면 수백만 달러가 든다는 게 업계의 상식처럼 굳어져 있습니다. 애니메이션 창작 도구를 만드는 Linum이 이 통념에 정면으로 도전하는 실험 결과를 내놓았습니다. GPU 비용 **총 21,038달러(약 2만 달러)**로 학습한 텍스트 기반 이미지 생성 모델 **Pyramid-JiT(P-JiT)**를 Apache 2.0 라이선스로 가중치까지 공개했어요.

이 모델은 차기 영상 모델 Linum v3로 가기 위한 **개념 검증(PoC)**입니다. 그런데 모델 크기에 비해 결과물이 좋아서 실험적으로 공개하기로 했다고 합니다. 어떤 아이디어가 이 정도 비용 절감을 가능하게 했는지 하나씩 살펴보겠습니다.

핵심 요약

총 비용 21,038달러: H100 8,415시간 (사전학습 20,348달러 + 후속학습 690달러)
VAE 없는 픽셀 공간 학습: 압축과 생성을 확산 모델 하나가 함께 학습
32×32 토큰 축소: 기존 Linum v2 대비 어텐션 문맥 길이를 4분의 1로 줄임
중간 블록 감독 신호: DiT 중간에서 저해상도 이미지를 예측해 수렴 속도 향상 (초기 실험 기준 11배 적은 샘플로 수렴)
모델 구성: 2B 규모 DiT + 4B 텍스트 인코더(Qwen3.5-4B), 512²에서 사전학습한 뒤 1024²로 확장
한계: 글자·포스터 생성 불가, 전신·다수 객체·여러 사람의 상호작용·세계 지식에 약함

왜 이런 실험을 했을까?

Linum의 목표는 누구나 애니메이션을 만들 수 있는 창작 도구입니다. 그런데 지금의 생성형 영상 모델은 이런 문제를 안고 있습니다.

느리고 비쌉니다
사용자가 결과를 다룰 제어 수단이 부족합니다
만들고 고치기를 반복해야 하는 창작 과정에서 이 두 가지 제약이 스토리텔링을 어렵게 합니다

그래서 Linum v3를 개발하면서 추론 비용을 크게 낮추는 것을 최우선으로 삼았습니다. 생성형 영상 모델은 보통 이미지 모델에서 출발해 사람·장소·사물을 먼저 익히고, 그다음에 동작과 움직임을 배웁니다. 그래서 이번 실험의 질문은 이것이었어요.

"5,000만~1억 샘플 규모 초기 실험에서 효과를 본 아키텍처가 약 6억 샘플 규모에서도 잘 작동할까?"

VAE를 없앤 이유: 어텐션 비용과의 싸움

병목은 O(n²) 어텐션

영상 학습과 추론에서 가장 큰 병목은 문맥 길이의 제곱에 비례하는 어텐션 연산입니다. 압축을 더 강하게 해서 토큰 수를 줄이면 학습 비용과 추론 비용이 함께 내려갑니다.

별도 VAE가 꼭 최선은 아니다

지금까지는 VAE로 이미지를 잠재 공간에 압축한 뒤 그 위에서 확산 모델을 돌리는 방식이 표준이었습니다. Linum은 여기에 몇 가지 문제를 짚습니다.

VAE 압축률에는 경험적인 한계가 있고, 더 높이려면 Minimax H3의 2B 디코더처럼 구조를 크게 바꿔야 합니다
VAE의 재구성 성능이 좋아져도 후단 확산 모델의 생성 품질은 오히려 나빠질 수 있습니다

그래서 P-JiT는 더 강력한 생성 모델인 확산 모델이 압축과 생성을 동시에 학습하도록 설계했습니다. 다른 모델이 만든 잠재 공간을 빌려 쓰지 않고, 자기 용도에 맞는 잠재 공간을 직접 학습하면 압축은 더 강해지고 학습은 더 빨라질 것이라는 가설입니다.

16×16 → 32×32

| 구분 | Linum v2 | P-JiT | |---|---|---| | 구조 | Wan VAE + DiT | VAE 없음, 픽셀 공간 DiT | | 토큰 축소 | 16×16 (FLUX-2 등과 동일) | 32×32 | | 효과 | — | 256² 이미지와 같은 토큰 수로 512² 이미지 표현 |

초기 실험에서 토큰 축소는 4배 강해졌고, 수렴에 필요한 샘플 수는 11분의 1로 줄었습니다. 핵심 변경은 세부 묘사 학습을 돕는 노이즈 스케줄과 DiT 본체 중간에 붙인 출력 헤드 두 가지입니다.

모델 구조 들여다보기

기본 구성

파라미터: 추론 시 2.18B, 학습 전용 헤드(약 0.24B)를 더하면 약 2.42B
흐름: 텍스트와 이미지가 각각 단일 모달리티 정제용 DiT 블록 2개를 거친 뒤, DiT 블록 22개로 이루어진 하나의 흐름으로 합쳐집니다
텍스트 인코더: 가중치를 고정한 Qwen3.5-4B의 전체 어텐션 계층(7, 15, 27번 층)에서 표현을 뽑아 이어 붙이고 투영합니다
이미지 패치: 32×32 패치(RGB 기준 3,072차원)를 1,024차원으로 압축합니다. 압축비는 3:1입니다
세부 수치: 은닉 차원 2,944, 128차원 어텐션 헤드 23개, FFN 내부 폭은 본체 5,376 / 정제부 2,816

초기 블록의 약한 기울기 문제를 푸는 방법

REPA, RAE-v2, Self-Flow 같은 연구는 모두 DiT 초기 블록에 기울기 신호가 약하게 전달되는 문제를 다룹니다. P-JiT는 픽셀 공간에서 학습한다는 점을 활용합니다. 픽셀 공간에서는 각 블록이 최종 이미지의 한 버전을 예측할 수 있으니, 중간에 정답을 직접 알려주는 방식이에요.

5번 블록: 2층 PixelREPA 모듈로 잠재 공간을 DINOv3-L에 정렬합니다. 은닉 토큰 20%를 가린 뒤 같은 이미지의 DINOv3-L 특징을 예측합니다
10번 블록 뒤: 4분의 1 해상도의 깨끗한 이미지를 예측합니다
16번 블록 뒤: 절반 해상도의 깨끗한 이미지를 예측합니다

이름에 '피라미드(Pyramid)'가 붙은 이유가 여기 있습니다. 저해상도에서 고해상도로 단계적으로 예측하는 구조죠.

총 손실은 중간·최종 출력의 MSE, PixelREPA 코사인 거리, 전체 해상도 이미지에 대한 지각 손실(LPIPS, P-DINO)을 합한 값입니다.

타임스텝 조건 주입

Z-Image를 따른 저랭크 AdaLN으로 각 블록에 타임스텝 조건을 주입합니다
타임스텝 임베딩을 256차원으로 줄이고, 블록별 확장 투영을 0으로 초기화해 각 블록이 항등 변환으로 시작하도록 했습니다

노이즈 커리큘럼과 고해상도 전환

큰 구조를 먼저, 세부는 나중에

첫 8,000만 샘플은 높은 노이즈에 집중해 기본 구조를 익힙니다
그 뒤 타임스텝 분포를 LogitNormal(0.8, 0.8) → LogitNormal(-0.2, 1.0)으로 넓혀 세부 묘사를 학습합니다
8,000만 샘플 시점에 검증 손실이 튀는 건 분포를 넓힌 직후 첫 평가이기 때문입니다

EMA 운영

구조와 세부를 어느 정도 익힌 1억 샘플부터 EMA를 시작합니다 (반감기 400만 샘플)
기울기 노름이 이동 중앙값의 10배를 넘으면 EMA 갱신을 250스텝 동안 멈춰 불안정한 업데이트를 걸러냅니다

512² → 1024²

512² EMA 가중치에서 새 옵티마이저로 다시 시작합니다
해상도가 높아지면 정보 중복도 늘어나기 때문에 노이즈 스케일을 2에서 4로 올렸습니다 (JiT 부록 A 근거)
RoPE는 512² 기준 격자를 유지하고 1024² 패치 인덱스에 ½을 곱해 좌표를 맞춥니다
전체 손실 값은 높아지지만, 이미지 품질은 빠르게 좋아졌다고 합니다

학습 규모와 비용 정리

학습 설정

데이터 규모: 512²에서 약 6억 1,000만 샘플, 1024²에서 약 5,000만 샘플 (총 약 6억 6,000만)
종횡비: 정사각형을 포함한 5개 종횡비를 20%씩, 이미지 크기 재조정 없이 사용
하드웨어: H100 64개 (8개 노드), 전역 배치 크기 2,048
정밀도/병렬화: bf16, FSDP2
학습 목표: x-prediction + v-loss
옵티마이저: 2차원 블록 행렬에는 Muon, 나머지에는 AdamW (Muon 업데이트 RMS를 AdamW에 맞춤)
학습률: 최대 2e-4로 올려 유지한 뒤 마지막 약 6,000만 샘플 동안 1 − √ 감쇠로 2.5e-5까지 내림
기타: 가중치 감쇠 0.1, 기울기 노름 1.0 제한, 분류기 없는 가이던스를 위해 캡션 5% 드롭

비용 (H100 시간당 2.50달러, 재시작 비용 포함)

| 단계 | H100 시간 | 비용 | |---|---|---| | 사전학습 | 8,139시간 | 20,348달러 | | 후속학습 | 276시간 | 690달러 | | 합계 | 8,415시간 | 21,038달러 |

확산 모델 후속학습, 아직 갈 길이 멀다

후속학습에는 FDFO(Finite Difference Flow Optimization) 변형을 쓰고, 이미지 평가 모델 JEV-Image와 사람의 평가를 함께 활용했습니다. Linum은 이 분야의 현주소를 이렇게 정리합니다.

모델·스케줄 증류는 많이 연구됐지만, 지난 18개월간 LLM 코딩 능력에서 나온 수준의 후속학습 성과가 확산 모델에서는 아직 나오지 않았습니다
DPO는 특정 미적 취향에 크게 과적합하는 경향이 있습니다
Flow-GRPO는 검증 가능한 보상 기반 강화학습(RLVR)을 확산 모델에 도입한 초기 작업이고, 이후 Diffusion NFT, FDFO 등이 뒤를 이었습니다
Linum은 FDFO를 지금까지 나온 것 중 확산 모델 후속학습을 가장 잘 정식화한 작업으로 평가합니다
자기회귀 LLM과 달리 믿을 만한 우도 추정치가 없다는 점이 기법 적용을 어렵게 만든다고 봅니다

후속학습의 전체 구성은 앞으로 따로 공개할 계획이라고 하니 기다려 볼 만합니다.

알아두어야 할 한계

글자는 못 씁니다

텍스트가 렌더링된 이미지를 학습 데이터에서 의도적으로 뺐습니다. 인포그래픽, PDF, 책 페이지, PowerPoint 슬라이드 같은 이미지는 애니메이션 도구라는 목표에 쓸모가 없다고 봤기 때문입니다. 그래서 글자나 포스터는 생성하지 못합니다.

작은 모델의 한계

2B DiT + 4B 텍스트 인코더라는 작은 구성 탓에 다음 영역에 약합니다.

전신 표현
자동차 6대나 군중처럼 개체가 많은 장면
여러 사람의 상호작용
세계 지식

이런 이미지도 가끔은 정확하게 나옵니다. Linum은 7B급 모델로 키우고 후속학습을 더하면 이 문제가 완화될 수 있다고 봅니다.

마무리: 이 실험이 의미 있는 이유

P-JiT는 최고 성능을 노린 모델이 아닙니다. 그래도 눈여겨볼 이유는 충분합니다.

비용: 개인 연구자나 작은 스타트업도 해볼 만한 2만 달러대 예산으로 쓸 만한 이미지 생성 모델을 만들 수 있다는 걸 보여줬습니다
구조: "VAE + 잠재 확산"이라는 표준 구조 대신 픽셀 공간에서 압축과 생성을 함께 학습하는 방식이 대규모에서도 통한다는 근거를 더했습니다
영상 모델로의 확장성: 어텐션 문맥을 4분의 1로 줄인 효과는 토큰 수가 훨씬 많은 영상에서 더 크게 나타납니다
개방성: Apache 2.0 라이선스로 가중치를 공개해 상업적 활용과 후속 연구가 자유롭습니다

생성형 AI의 비용 장벽을 낮추는 시도에 관심이 있다면, P-JiT 부록과 데이터 필터링 기록까지 꼭 읽어보시길 권합니다. Linum v3가 영상에서 어떤 결과를 보여줄지도 기대됩니다.

출처: GeekNews

원문: https://news.hada.io/topic?id=36694&utm_source=geeknews&utm_medium=rss&utm_campaign=news