Beyond the Estimate: Wiener-Informed Decoding for Easier Training of Pixel-Space Flow Models

Learning velocity fields directly in pixel space is difficult and slow. We show that coarse second-order image statistics can ease this task: frozen Wiener estimates improve velocity learning, with gains extending beyond the linear predictions they explicitly provide. Moreover, we observe similar but complementary gains from frozen bases and pixel decoding. These observations motivate Wiener-informed Decoding (WiD), which uses statistical knowledge both for direct prediction and to organize decoder computation without modifying native blocks. WiD combines a trainable Wiener prediction path with coefficient interfaces that inject features into the native decoder and return its predicted coefficient corrections, with these components initialized with data statistics. All introduced parameters, including the spectrum and coordinate transforms, are trained jointly under the original pixel-space velocity objective. Across model scales and decoder architectures, WiD accelerates training and improves generation quality with modest parameter overhead. It also yields smoother gradient trajectories and remains effective with representation alignment. WiD improves class-conditional and text-to-image generation across different resolutions. On ImageNet 256 x 256, PixelDiT-XL with WiD reaches FID 1.53 after 500 epochs, achieving comparable quality to PixelDiT with 37.5% fewer training epochs. Statistical denoisers contribute beyond the estimate: their computational structure can also ease pixel-space training.

한 줄 요약

Wiener-informed Decoding(WiD)는 픽셀 공간 흐름 모델의 학습을 가속화하고 생성 품질을 향상시키는 통계 기반 디코딩 방법이다.

핵심 기여도

**Wiener 추정 기반의 통계적 예측 경로**를 도입하여 픽셀 공간 속도 학습을 용이하게 함.
**Coefficient interface**를 통해 통계적 좌표를 디코더에 통합, 기존 디코더 블록을 수정하지 않음.
**PixelDiT-XL**에서 WiD 적용 시 500 에포크에 FID 1.53 달성, 기존 PixelDiT 대비 37.5% 적은 에포크로 동일 품질 유지.
**Parameter overhead는 적은데도** 학습 안정성과 효율성, 생성 품질 모두 개선됨.

핵심 아이디어

픽셀 공간에서 속도 필드를 학습하는 것은 고차원성과 복잡한 이미지 구조로 인해 어려운 문제이다. 이에 저자들은 **Wiener 추정**을 활용한 새로운 접근법을 제안한다. Wiener 추정은 이미지의 2차 통계량을 기반으로 선형적인 노이즈 제거를 제공하지만, 저자들은 이 추정의 **계산 구조 자체가 학습을 도울 수 있음을 발견**했다.

Wiener-informed Decoding(WiD)은 **통계적 좌표 변환**(analysis/synthesis transform)을 디코더에 통합하여, **직접적인 속도 예측과 함께 통계적 예측을 보완하는 방식**으로 학습을 진행한다. 이는 기존 디코더의 입력 임베딩에 통계적 정보를 추가하고, **픽셀 공간과 계수 공간의 예측을 결합**하여 최종 속도를 예측하는 방식이다.

**patch PCA 기반의 Wiener 경로**는 노이즈 패치를 통계적 좌표로 변환하고, 이 좌표를 기반으로 속도 예측을 생성한다. 이는 디코더의 학습을 안정화하고, 초기 학습 단계에서의 성능 저하를 완화한다.

기술적 접근법

**Wiener 경로**: `patch PCA` 기반의 분석 변환 `U_enc`와 합성 변환 `U_dec`을 사용하여 노이즈 패치를 통계적 좌표로 변환하고, Wiener 이득 `a(t; λ)`를 적용.
**Coefficient interface**: 통계적 좌표는 디코더의 입력 임베딩 `ℰ`에 추가되어, `W_F`를 통해 픽셀 수준 특징으로 변환됨.
**Decoder output**: 디코더는 픽셀 공간 속도 `v`와 계수 공간 보정 `Δc`를 동시에 예측.
**Loss**: 두 예측의 합 `v̂_θ`가 속도 손실 함수(velocity objective)로 감독됨.
**Parameter training**: `U_enc`, `U_dec`, `λ`, `W_F`, `W_Δ`, `b_Δ` 등 모든 추가 파라미터는 기존 속도 손실과 함께 학습됨.

주요 결과

**ImageNet 256×256** 데이터셋에서 PixelDiT-XL에 WiD 적용 시 500 에포크에 **FID 1.53** 달성.
기존 PixelDiT 모델이 동일 품질을 달성하기 위해 **37.5% 더 많은 에포크**가 필요.
**Gradient trajectory가 더 매끄럽고**, 학습 안정성이 향상됨.
**Representation alignment** 상황에서도 WiD의 효과 유지됨.
**Class-conditional 및 text-to-image 생성**에서 다양한 해상도에서 성능 개선.

의의 및 한계

WiD는 픽셀 공간 흐름 모델의 학습 과정에서 **통계적 구조를 활용한 새로운 디코딩 전략**을 제시한다. 기존 연구는 주로 예측 공간의 파라미터화나 아키텍처 변경에 집중했으나, WiD는 **통계적 계산 구조 자체를 학습 프로세스에 통합**함으로써 학습 효율성을 높인다. 이는 **고차원 픽셀 공간에서의 학습 난이도를 완화**시키는 데 기여하며, 향후 픽셀 공간 생성 모델의 설계에 새로운 가능성을 제시한다.

하지만, WiD는 **기존 디코더 아키텍처에 의존**하며, 통계적 기반(`patch PCA`)이 이미지의 복잡한 비선형 구조를 완전히 포착하지 못할 수 있다. 또한, **데이터셋의 통계가 충분히 잘 추정되지 않으면 성능 저하가 발생할 수 있음**이 명시되지 않음.

실용적 활용

WiD는 **고해상도 이미지 생성**, **클래스 조건부 생성**, **텍스트-이미지 생성** 등 다양한 생성 모델 태스크에 적용 가능하다. 특히, **사전 학습된 오토인코더 없이도 픽셀 공간에서 직접 생성하는 모델**에 유용하며, **학습 시간 단축과 생성 품질 향상**을 동시에 달성할 수 있다. 이는 **실시간 생성, 대규모 데이터셋 학습, 리소스 제한 환경**에서 실용적 가치가 높다.