공유 블록 반복으로 파라미터를 고정한 이미지 생성 모델 Looped-DiT

무엇인가

텍스트-이미지 모델을 개선하는 기존 방식은 모델 크기를 키우거나 denoising step 수를 늘리는 것이었다. Qwen-Image, FLUX.2 같은 모델이 수십억 파라미터로 커지면서 배포 비용이 커졌다. 이 논문은 세 번째 축을 제안한다. 각 denoising step 안에서 공유 Transformer 블록을 N번 반복 실행해, 파라미터 수와 시퀀스 길이를 고정한 채 계산 깊이만 늘리는 방식이다. 저자들은 이미지 생성이 프롬프트에 암시된 시각적 내용을 추론하고, 서로 얽힌 제약을 해소하고, 생성이 진행되며 생긴 불일치를 찾아 고치는 과정을 요구한다고 본다. 이런 과정이 본질적으로 시각적이므로, 텍스트 추론 토큰 없이 hidden representation 안에서 직접 반복 정제할 수 있다는 것이 이 접근의 동기다. 중간 loop의 예측을 같은 target 이미지로 감독할 수 있어 diffusion 학습 파이프라인에도 자연스럽게 얹힌다.

어떻게 동작하나

문제는 순진한 looping이 통하지 않는다는 점이다. 논문의 초기 실험에서 loop depth가 얕을 때는 non-looped 베이스라인보다 성능이 낮게 머물렀고, loop를 더하면 성능이 포화되다가 학습 시 사용한 loop depth를 넘어서면 오히려 떨어졌다. 저자들은 같은 변환을 반복 적용하면 중복 업데이트가 image-token 표현의 정보를 덮어쓰거나 약화시킨다고 가설을 세웠다. 이를 확인하기 위해 각 loop depth에서 hidden state로부터 image token의 2D patch-grid 좌표를 예측하는 ridge regression probe를 학습시켰다. 그 결과 결정계수 R²가 1번째 loop 후 0.865에서 8번째 loop 후 0.562로, 0.303만큼 떨어졌다. 공간 위치 정보가 loop를 거듭할수록 선형적으로 디코딩하기 어려워진 것이다. 결론은 두 가지 요구사항으로 정리된다. 중간 loop 예측에 대한 실질적인 감독과, 표현이 변해가는 동안 attention 업데이트 강도를 적응적으로 조절하는 능력이다.

무엇과 다른가

Looped-DiT는 pixel-space MMDiT인 MiniT2I를 백본으로 삼는다. 네트워크를 pre-loop 단계 A, looped 단계 B, post-loop 단계 C로 나누고, h(0)=A(h_input), h(r)=B(h(r-1)) (r=1,…,N), x̂0=C(h(N))으로 계산한다. B의 파라미터는 모든 반복에서 공유되므로 N을 키워도 파라미터는 늘지 않는다. 구현에서는 17개 MMDiT 블록을 6/5/6으로 쪼개고 가운데 5개 블록을 N=4회 반복한다. 이 looping은 각 denoising step 내부에서 sampler가 진행되기 전에 일어나므로, loop depth와 denoising step 수를 서로 독립적으로 조절할 수 있다. 패치 크기에 따라 B/16과 B/32 두 변형을 학습했다.

어떻게 쓰나

첫 번째 구성 요소는 Deep Supervision이다. 최종 예측만 감독하면 앞선 loop 상태까지 도달하는 데 모든 후속 반복을 거쳐야 해서 최적화 경로가 길고 간접적이 된다. 그래서 각 loop depth n의 hidden state를 공유 post-loop 단계 C로 디코딩해 x̂0(n)을 얻고, flow-matching 목적함수를 clean-image 예측 형태로 모든 loop에 적용한다. 손실은 ℓ_n = E[‖x̂0(n)−x0‖²/(d_x·c(t)²)]이고 c(t)=max{1−t, τ}로 t가 1에 가까워질 때 가중치가 발산하는 것을 막는다. 전체 목적함수는 각 loop 손실의 가중합 Σ w_n ℓ_n이다. 중간 예측은 학습 중에만 디코딩되므로 추가 파라미터나 추론 오버헤드가 없다. 두 번째는 Self-Modulating Attention이다. 표준 softmax attention은 source token의 상대적 기여도만 조절할 뿐 업데이트의 세기는 제어하지 못한다. 저자들은 Gated Attention, 즉 각 attention head 출력을 학습된 스칼라 게이트 σ(w_g^T u + b_g)로 명시적으로 스케일하는 방식과, Exclusive Self Attention(XSA), 즉 토큰 자기 자신의 value 방향 성분을 제거하는 파라미터 없는 상태 의존 projection(I − v̂v̂^T)을 비교한다. 두 변조 기법 모두 looped 단계 B 안에만 적용한다.

전제와 한계

실험은 DPG-Bench, PRISM, T2I-CoReBench, SpatialGenEval, GenEval, TIIF-Short에서 이뤄졌다. B/16은 DPG-Bench, PRISM, T2I-CoReBench, SpatialGenEval, TIIF-Short에서 최고 성능을 냈고, 차선 모델인 InternVL-U보다 평균 점수가 2.5점 높으면서 파라미터는 6.5배 적다. 260M 파라미터 looped 모델이 6.5배 큰 모델을 여러 벤치마크에서 앞서고 추론 연산은 4.9배 적다는 것이 논문의 대표 수치다. 추론 연산과 지연 시간 양쪽에서 Pareto frontier에 올랐다. 공정 비교에서는 B/32가 파라미터를 맞춘 베이스라인보다 평균 4.0점 높았고, 학습 연산을 맞춘 조건에서 59.1 대 58.1로 더 깊은 non-looped 모델을 앞섰다. 추론 예산을 고정하고 25 step에서 50 step 구간을 비교하면, denoising step을 늘리는 것보다 loop depth를 늘리는 쪽이 일관되게 더 높은 성능을 냈다. 텍스트 CoT와의 비교에서는 관계·절차·공간 제약을 동시에 만족해야 하는 하위 과제에서 looping의 이득이 컸고, 예시와 규칙으로부터 암시적 시각 내용을 추론해야 하는 하위 과제에서는 CoT의 이득이 컸다. 둘을 결합하면 전체 점수가 가장 높았다.

Ablation 결과도 구체적이다. Deep Supervision 가중치 스킴은 Exponential과 Final+Mean 둘 다 final-loop-only보다 좋았고 Final+Mean이 더 우수했다. Self-Modulating Attention은 Gated와 XSA 모두 변조 없는 attention보다 좋았고 XSA의 이득이 더 컸으며, Final+Mean 감독과 XSA를 결합한 구성이 최고 평균 점수를 냈다. Deep Supervision은 1~8개 inference loop 전반에서 성능을 끌어올리고 early exit에 대한 페널티를 크게 줄였으며, 학습에 쓴 4 loop를 넘어서도 성능을 유지했다. loop로 학습한 모델은 1회만 돌려도 표준 non-looped 베이스라인을 앞섰다. XSA의 이득은 looped 모델에서 +1.6이었지만 파라미터를 맞춘 non-looped 베이스라인에서는 −0.1, 연산을 맞춘 베이스라인에서는 +0.7에 그쳤다. 변조 없는 attention은 loop가 늘수록 상대 업데이트 노름이 가장 크고 위치 디코딩 가능성(R²) 하락이 가장 컸으며, XSA는 업데이트가 가장 작고 위치 디코딩 가능성이 가장 높았다. 이 순서는 벤치마크 성능 순서와 일치한다. 논문은 또한 4개의 learnable query를 쓰는 cross-attention과 MLP head로 각 loop 이후 계속할 가치가 있는지를 예측하는 경량 gating network를 붙인 adaptive looping을 제안한다. 학습 타깃은 y_r = max{0, max_k (ℓ_r−ℓ_k)/(k−r)}이고, 추론 시 임계값 λ로 평균 loop 수를 재학습 없이 조절한다. Adaptive looping은 평균 loop 수를 맞춘 fixed-loop 추론보다 우수했고, 1.25 loop 조건에서 성능 하락을 절반으로 줄였으며 2.70 loop에서 성능 plateau에 더 일찍 도달했다.

개발자 관점에서 이 논문이 유용한 지점은 추론 예산과 파라미터 예산이 모두 묶여 있는 배포 환경이다. GPU 메모리 때문에 모델 크기를 못 키우지만 지연 시간은 어느 정도 허용되는 경우, 같은 체크포인트에서 loop depth를 늘려 품질을 사는 선택지가 생긴다. loop depth와 denoising step이 독립적이라는 점도 실무적으로 중요하다. 같은 가중치로 두 축을 조합해 예산을 맞출 수 있고, adaptive looping을 쓰면 프롬프트 난이도에 따라 평균 loop 수를 낮출 수 있다. 다만 이 결과가 MiniT2I라는 최소 구성의 pixel-space MMDiT, 17개 블록 중 5개를 4회 반복하는 설정에서 나왔다는 점은 확인이 필요하다. 다른 아키텍처 규모나 latent-space 모델에서 loop 구간을 어디에 두고 몇 번 돌려야 하는지, Deep Supervision 가중치 스킴이 그대로 통하는지는 이 논문이 답하지 않는다. 코드는 공개돼 있다.

논문은 별도의 한계 절을 두지 않았지만 본문에 전제가 드러나 있다. 저자들은 순진한 looping이 신뢰할 수 없고 공간 정보를 점진적으로 훼손한다는 것을 자기 실험으로 먼저 보였고, 그래서 중간 감독과 적응적 조절을 필수 조건으로 못박는다. 실험은 B/16과 B/32 두 변형에 한정되며, 주요 결과는 B/16의 6개 데이터셋 평균, 분석과 ablation은 B/32의 reasoning 관련 4개 데이터셋 평균으로 보고된다. "latent reasoning"이라는 표현도 저자들 스스로 "suggestive of"라고 표현한 정성적 증거 수준이며, 별도의 추론 벤치마크로 정량화하지는 않았다. adaptive looping의 임계값 λ는 재학습 없이 조절되지만 gating network 자체는 별도로 학습해야 한다.