초저비용 이미지 생성 AI P-JiT 공개됐다

연구자 2명이 꾸린 AI 연구 조직 리넘(Linum)이 이미지 생성 모델 피라미드-JiT(Pyramid-JiT, P-JiT)를 오픈 웨이트로 공개했다. 파라미터 20억개 규모인 이 모델은 학습 비용을 2만1,000달러(2,940만원대)로 낮췄다.

P-JiT는 이미지 생성 AI가 안고 있는 두 가지 과제를 겨냥했다. 생성에 시간과 비용이 많이 든다는 점과 제어가 어렵다는 점이다. 기존 이미지 생성 AI는 VAE로 이미지를 잠재 표현으로 바꾼 뒤 다시 이미지로 디코딩하는 과정을 거친다. P-JiT는 잠재 표현으로 바꾸지 않고 이미지를 픽셀 데이터 그대로 처리한다.

◇ 토큰 압축해 학습 표본 11분의 1로=기존 생성 모델은 어텐션 연산량이 입력 시퀀스 길이의 제곱에 비례한다. 동영상 생성 모델에서는 이 때문에 학습과 추론에 막대한 시간과 비용이 들었다. P-JiT는 토큰을 적극적으로 압축해 어텐션 컨텍스트 창을 줄이고 연산 비용을 크게 낮췄다.

이전 모델 리넘 v2는 512² 이미지를 16×16 단위로 줄여 토큰으로 표현했다. P-JiT는 이를 32×32 단위로 늘렸다. 어텐션 컨텍스트 창을 4분의 1로 줄인 셈이다. 그 결과 리넘 v2보다 11분의 1 수준의 학습 표본으로 학습이 수렴했다.

P-JiT는 픽셀 공간 모델에 맞춘 노이즈 스케줄을 적용했다. 디퓨전 트랜스포머(DiT) 앞쪽 블록의 기울기 신호를 키우는 판독(readout) 장치도 넣었다. 네트워크 초반에 저해상도 이미지를 미리 예측하게 해 수렴을 앞당기는 장치다.

텍스트와 이미지는 각각 전용 리파이너 DiT 블록 2개를 거친 뒤 동일한 DiT 블록 22개로 이뤄진 하나의 흐름으로 합쳐진다. 텍스트 인코더로는 파라미터를 고정한 큐원(Qwen)3.5-4B 모델을 쓰고 7·15·27번째 레이어 출력을 활용한다. 이미지는 32×32 패치로 나눈 뒤 1,024차원 잠재 공간에 3대 1 비율로 압축해 이미지 리파이너 블록에 넣는다. 5번 블록에는 잠재 공간을 DINOv3-L에 맞추는 2층짜리 PixelREPA 판독 회로를 달았다. 10번과 16번 블록 뒤에는 선형 판독 헤드 2개를 두고 각각 4분의 1 해상도와 2분의 1 해상도로 깨끗한 이미지를 예측하게 했다.

학습에는 MSE와 PixelREPA 기반 코사인 거리, LPIPS와 P-DINO 기반 지각 손실 등 여러 손실 함수를 함께 썼다. 학습은 단계적으로 진행했다. 512² 이미지로 표본 6억6,000만개를 학습한 뒤 1,024² 이미지로 표본 5,000만개를 학습하는 커리큘럼 학습 방식이다.

◇ 글자 든 이미지는 못 만들어=P-JiT 학습 데이터에는 인포그래픽이나 PDF, 책 페이지가 들어가지 않았다. 그래서 글자가 들어간 이미지나 포스터, 인포그래픽은 생성하지 못한다. 비교적 작은 모델이라 전신 묘사와 여러 개체 묘사, 인물 간 상호작용, 폭넓은 세계 지식에서도 한계가 있다. 리넘은 더 큰 모델로 학습하거나 사후 학습을 더하면 이런 문제를 줄일 수 있다고 봤다. 주사위 눈 배치를 몰라 주사위를 틀리게 그린 사례도 실패 예시로 소개했다. 관련 내용은 이곳에서 확인할 수 있다.