픽셀 확산 트랜스포머의 균일 정제를 깨서 전역 구조와 고주파 세부를 분리한다
Persistence Forcing: Exploiting Feature Specialization in Pixel-Space Diffusion
무엇인가
픽셀 공간 확산 트랜스포머(DiT)는 사전학습된 토크나이저 없이 고차원 픽셀을 직접 다루기 때문에, 유용한 내부 표현을 스스로 조직하는 일이 성능의 핵심 병목이 된다. 그런데 기존 DiT는 모든 은닉 특징 차원이 네트워크 깊이 전체를 통과하며 동일한 횟수로 변환되는 균일 정제(uniform refinement) 패턴을 유지한다. 자연 이미지는 전역 구조는 압축적으로 표현할 수 있지만 국소 질감과 고주파 세부는 훨씬 풍부한 표현을 요구한다는 점에서, 이 논문은 모든 특징이 같은 양의 정제를 받아야 하는지를 다시 묻는다.
어떻게 동작하나
제안하는 첫 단계는 이질적 정제(heterogeneous refinement)다. D차원 은닉 표현을 유지하되 ℓ번째 트랜스포머 블록이 앞쪽 dℓ ≤ D개 차원만 실제로 변환하고 나머지는 그대로 우회시킨다. 각 특징 i의 정제 예산은 R_i = Σ 1[i ≤ dℓ]로 정의되며, 균일 정제에서는 모든 차원이 R_i = L이지만 dℓ을 깊이에 따라 바꾸면 차원마다 서로 다른 정제 이력이 생긴다. 저자들은 활성 폭을 점점 좁혔다가 다시 넓히는 수축–확장 스케줄을 사용하고, 백본 용량 변화와 효과를 분리하기 위해 Σ dℓ² = L d²를 유지해 기존 JiT와 어텐션·MLP 파라미터 예산을 근사적으로 맞춘다.
무엇과 다른가
이 최소한의 개입만으로 순서 있는 특징 분화가 창발한다. 정제를 적게 받은 그룹은 일관된 객체 외형과 전역 이미지 구조를 이미 보존하고, 정제를 많이 받을수록 기여가 질감·경계·잔차 같은 국소 고주파 콘텐츠로 이동한다. 저자들은 전자를 지속 특징(persistent), 후자를 활성 특징(active)이라 부른다. 중요한 것은 이질적 정제만으로는 생성 성능이 좋아지지 않는다는 점이다. 파라미터를 맞춘 이질적 정제 모델은 같은 조건의 바닐라 JiT를 능가하지 못하고 FID가 오히려 나빠질 수 있다. 남은 기회는 두 역할이 어떻게 상호작용하느냐에 있다.
어떻게 쓰나
Persistence Forcing(PerF)은 지속 표현을 활성 정제의 구조적 문맥으로 명시적으로 연결한다. ℓ번째 블록에서 활성 특징 hℓᵃ는 트랜스포머 블록이 갱신하는 반면, 우회하는 지속 특징 hℓᵖ는 변환되지 않고 추가 변조 조건으로 쓰인다. 구체적으로 표준 타임스텝·클래스 변조에 지속 표현에서 유도한 토큰별 변조를 더한다: mℓ = AdaLNℓ(t, y) + Pℓ(Norm(hℓᵖ)). Pℓ은 랭크 r의 저랭크 병목으로 구현해(PerF-B/L은 128, PerF-H는 192) 추가 파라미터를 5% 미만으로 유지하면서, 지속 정보가 매 활성 정제 단계에 영향을 주되 자기 자신은 반복 변환되지 않게 한다. 여기에 더해 학습 중 지속 조건 경로를 확률적으로 꺼서(클래스 조건과 지속 조건을 각각 0.1 확률로 드롭) 같은 모델이 조건 유무 양쪽 예측을 낼 수 있게 하고, 샘플링 시 Persistence Guidance(PG)를 만든다. CFG가 x_uᵖ + w_c(x_cᵖ − x_uᵖ)로 외부 의미 조건을 대비시키는 것과 대칭적으로, PG는 x_c^∅ + w_p(x_cᵖ − x_c^∅)로 지속 조건의 유무만 다른 두 예측을 대비시켜 내부에서 형성된 구조적 문맥을 강화한다. 두 가이던스는 x_uᵖ + w_c(x_cᵖ − x_uᵖ) + (w_p − 1)(x_cᵖ − x_c^∅)로 결합된다.
전제와 한계
실험은 ImageNet 256×256과 512×512 클래스 조건부 생성에서 JiT 아키텍처를 기반으로 수행했다. 주요 실험에서는 지속 특징에 REPA를 가중치 0.1로 적용했다. 256×256에서 PerF는 모든 스케일에서 JiT를 개선해 FID를 B/16 3.66→2.81, L/16 2.36→1.91, H/16 1.86→1.63으로 낮췄고, 상대 감소폭은 각각 약 23%, 19%, 12%다. PerF-H는 Inception Score도 303.4→324.5로 올렸고, 세 스케일 모두 recall이 상승하는 동시에 precision은 JiT와 비슷한 수준을 유지해 다양성을 희생하지 않았음을 보였다. 512×512에서는 PerF-H/32가 FID 1.94→1.76, IS 309.1→335.3, recall 0.61→0.64를 기록했다.
절제 실험은 REPA 없이 100 에폭 학습한 PerF-B 기준으로 이뤄졌다(따라서 절대 FID 값이 본 실험보다 훨씬 높다). 수축–확장 스케줄을 네트워크 깊이에 따라 이동시켜도 FID는 40.35~41.35 범위로 거의 변하지 않았고, 같은 학습 단계의 균일 JiT 42.1보다 모두 우수했다. 병목이 늦게 오는 스케줄이 근소하게 좋아 Late Bottleneck을 기본으로 채택했다. 구성 요소별로는 이질적 백본만으로는 CFG 없이 42.71, CFG 사용 시 8.02였는데, P-to-A 조건화를 넣으면 40.35와 7.19로 좋아지고 PG까지 더하면 24.92와 5.81로 가장 큰 추가 개선이 나온다. 저랭크 병목 랭크를 128에서 256으로 늘리면 PG 없이는 40.35→40.06에 그치면서 파라미터가 5.7M→11.5M로 두 배가 되지만, PG를 강하게 걸면 랭크 확대 효과가 커진다. PG는 CFG 스케일을 따로 튜닝한 뒤에도 폭넓게 이득을 주며 최적점은 w_p ≈ 3 부근이고, 전체 디노이징 구간 [0,1]에 적용할 때 가장 좋았다.
실무 관점에서 이 논문의 매력은 백본을 갈아엎지 않는다는 점이다. 은닉 폭을 층별로 줄였다 늘리는 스케줄과 저랭크 변조 브랜치, 그리고 학습 시 조건 드롭 확률만 추가하면 되고 추가 파라미터는 5% 미만이다. 샘플링 단계에서는 w_p와 w_c를 독립적으로 스윕해 볼 수 있고, PG를 전체 타임스텝에 걸어도 해가 되지 않는다는 결과는 구현 부담을 낮춘다. 다만 절제 실험의 FID 절대값이 본 실험과 크게 다른 만큼, 자신의 학습 예산과 데이터 규모에서 스케줄·랭크·가이던스 조합을 다시 확인해야 한다.
논문이 명시적으로 밝힌 한계 절은 원문에 별도로 제시되지 않았다. 다만 본문에서 확인되는 전제는 분명하다. 이질적 정제만으로는 성능이 오르지 않고 P-to-A 조건화와 PG가 함께 있어야 효과가 나며, PG는 w_p가 너무 커지면 오히려 성능이 떨어진다. 병목 랭크는 품질과 파라미터 오버헤드 사이의 절충으로 기본값 128을 쓰며, 개선폭은 백본이 커질수록 줄어든다(23%→19%→12%). 또한 이질적 정제 연구는 JiT라는 단순 균일 폭 픽셀 DiT 위에서, ImageNet 클래스 조건부 생성이라는 단일 설정에서 검증됐다.