Linum AI, VAE 없는 픽셀 공간 모델 Pyramid-JiT 공개… 학습 효율 11.3배 향상

Linum AI가 VAE 없이 픽셀 공간에서 작동하는 텍스트-이미지 생성 모델 **Pyramid-JiT (P-JiT)**을 공개했습니다. 이 모델은 LLM과 유사한 디코더 전용 구조를 채택해 기존 인코더-디코더 설계의 확장성 한계를 해결하며, DiT 트렁크를 통해 여러 해상도에서 타겟 이미지를 예측합니다.

성능 및 효율성

P-JiT은 이전 모델 대비 학습 효율성과 이미지 품질 지표에서 큰 개선을 보였습니다.

• FD-DINOv2 점수: P-JiT은 27.9를 기록해 JiT-DDT(39.1)를 크게 앞섰고, 실제 이미지 간 기준선인 25.4에 근접했습니다. Latent diffusion 방식인 Linum v2는 79.9점을 받았습니다.
• 학습 효율성: P-JiT은 Linum v2의 최종 FD-DINOv2 점수에 도달하기 위해 샘플 수를 11.3배, GPU 시간을 4.3배(픽셀 수 4배 기준) 줄였습니다.
• 토큰 감소: Linum v2보다 더 공격적인 토큰 감소(512x512 이미지당 256 토큰)를 적용했음에도 눈, 주름, 털 등 디테일이 더 자연스러운 고품질 이미지를 생성합니다.

아키텍처 및 방법론

핵심 혁신은 기존 인코더를 K-Readouts로 대체한 것입니다. DiT 트렁크의 다양한 깊이에 출력 헤드를 배치해 128x128, 256x256, 512x512 등 여러 해상도에서 이미지를 예측하는 멀티 스케일 지도 학습을 수행하며, 이를 통해 VAE를 제거했습니다.

• 게이티드 잔차: 최신 LLM 연구에서 영감을 받은 Gated Residuals를 실험했으나, AdaLN + tanh + post-norm을 사용한 기본 아키텍처가 디테일 학습에 더 강력한 바운드를 제공해 기본 정규화 전략을 유지했습니다.
• 타임스텝 조건화: Flow matching 처리를 위해 Shift-based 대신 Scale-based 타임스텝 조건화를 사용했으며, 구조와 디테일 제어에 더 효과적이었습니다.
• 텍스트 인코더: 동결된 Qwen3.5-4B의 hidden states를 DiT 공간으로 투영해 사용합니다.

프롬프트 준수 벤치마크

P-JiT은 분포 매칭(FD-DINOv2)에서는 우수하지만, 카운팅과 색상 정확도 등 프롬프트 준수 작업에서는 아직 Linum v2보다 뒤처져 해당 영역에서 추가 학습이 필요함을 시사합니다.

• GenEval: Linum v2(0.855) > P-JiT(0.817) > JiT-DDT(0.808).
• 카운팅 작업: Linum v2는 0.80, P-JiT은 0.56을 기록했습니다.
• 색상 작업: Linum v2는 0.95, P-JiT은 0.90을 기록했습니다.

코드와 가중치는 연구 아티팩트로 Apache 2.0 라이선스로 공개되었습니다. Linum AI는 P-JiT의 학습을 계속 진행하며 Linum v3로 향하는 과정에서 업데이트된 체크포인트를 릴리스할 계획입니다.