Linum AI, $21,038로 학습한 2B 픽셀 스페이스 텍스트-이미지 모델 P-JiT 공개

• 2인 연구소 Linum AI가 총 $21,038의 비용으로 처음부터 학습한 2B 파라미터 픽셀 스페이스 텍스트-이미지 모델 P-JiT(Pyramid-JiT)를 공개함
• VAE(Variational Autoencoder)를 제거하여 경험적 압축 한계가 있는 VAE보다 더 공격적인 압축이 가능한 목적형 잠재 공간 학습 가능
• 32x32 토큰 감소를 통해 512² 이미지를 기존 256² 이미지와 동일한 토큰으로 표현하여 Linum v2 대비 어텐션 컨텍스트 윈도우를 4배 축소함
• 64개 H100 GPU를 8,415시간 동안 사용하여 학습했으며, 사전 학습 비용은 $20,348, FDFO 변형을 활용한 사후 학습 비용은 $690임

Architecture & Mechanism

• 텍스트와 이미지 스트림을 병합하는 22개의 동일한 블록을 가진 DiT 구조 사용, 은닉 크기는 2,944(23 헤드 x 128)임
• 텍스트는 동결된 Qwen3.5-4B 인코더(7, 15, 27층)와 단일 모달리티 리파이너 블록 2개로 처리됨
• 이미지 패치는 3,072차원 RGB에서 3:1 압축 비율로 1,024차원으로 압축된 후 리파이너 블록 2개로 처리됨
• 샘플링 시 총 파라미터는 2.18B, 학습 시 PixelREPA와 선형 리드아웃 헤드 등 학습 전용 헤드를 포함해 약 2.42B임

Training Regime

• 512²(610.27M 샘플)과 1024²(49.97M 샘플) 두 단계로 학습하여 총 약 660M 샘플 사용
• 2D 블록 행렬에는 Muon 옵티마이저, 나머지는 AdamW를 사용하며 피크 학습률은 2e-4임
• Linum v2보다 11배 적은 샘플로 수렴함
• 하드웨어 구성은 bf16과 FSDP2를 사용하는 8노드 64개 H100 GPU임

Limitations & Scope

• 렌더링된 텍스트로 학습되지 않아 텍스트가 포함된 인포그래픽, PDF, 포스터 생성 불가
• 2B 파라미터 규모로 인해 전체 신체 자세, 다수 엔티티(예: 자동차 6대), 다인 상호작용 처리에 어려움
• 추론 비용 절감 및 애니메이션 제어 개선을 목표로 하는 향후 생성형 비디오 모델의 개념 증명(Proof-of-Concept)으로 개발됨