양방향 정규화 흐름: 데이터와 잡음 사이의 자유로운 변환

초록

BiFlow는 전통적인 정규화 흐름(NF)의 전방 모델은 그대로 두고, 역방향 모델을 별도로 학습한다. 역 모델은 정확한 해석적 역함수를 요구하지 않으며, 자유로운 아키텍처와 손실 설계가 가능해 샘플링 속도가 크게 향상되고 이미지 품질도 개선된다.

상세 분석

본 논문은 정규화 흐름(NF)의 핵심 제약인 “전방 모델의 명시적 가역성”을 완화하고, 역방향 과정을 학습 가능한 별도 모델로 전환한다는 혁신적인 아이디어를 제시한다. 기존 NF는 Jacobian determinant를 계산 가능하도록 설계된 가역 변환들의 연쇄로 구성돼야 했으며, 이는 복잡한 구조(예: coupling layer, autoregressive flow)로 제한을 받았다. 특히 최근 TARFlow와 같은 Transformer 기반 자동회귀 흐름은 강력한 표현력을 제공하지만, 인퍼런스 시 causal masking에 의한 순차적 연산이 필수라 샘플링 속도가 수천 배 느려지는 병목을 안고 있었다.

BiFlow는 두 단계 학습을 도입한다. 1) 전방 모델 Fθ는 기존 NF와 동일하게 최대우도(MLE)로 학습한다. 2) 전방 모델을 고정한 뒤, 역방향 모델 Gϕ를 학습한다. 여기서 역 모델은 가역성을 강제하지 않으므로, 비인과적(bidirectional) Transformer, U‑Net, ConvNeXt 등 자유로운 아키텍처를 사용할 수 있다. 논문은 역 모델 학습을 위한 세 가지 전략을 비교한다.

• Naïve Distillation: 최종 출력만을 L2 등 거리 손실로 맞추는 가장 단순한 방식으로, 역 모델이 전체 변환을 한 번에 학습해야 하므로 학습이 불안정하고 성능이 저조하다.
• Hidden Distillation: 전방 흐름의 중간 hidden state {xi}와 역 흐름의 중간 state {hi}를 직접 매칭한다. 이는 단계별 역변환을 강제하지만, 각 중간 state가 입력 차원과 동일해야 한다는 제약으로 인해 역 모델 설계가 경직된다.
• Hidden Alignment(제안 방식): 중간 state를 그대로 매칭하는 대신, 역 모델의 중간 출력 hi에 learnable projection head φi를 적용해 φi(hi)와 전방의 xi를 정렬한다. 이렇게 하면 역 모델이 자유롭게 표현 공간을 탐색하면서도 전방 trajectory 전체를 활용한 풍부한 지도 신호를 받을 수 있다.

또한, 기존 TARFlow는 잡음이 섞인 중간 표현 ˜x를 생성한 뒤 별도의 score‑based denoising 단계를 수행해 품질을 높였지만, 이는 추가 연산 비용을 초래한다. BiFlow는 이 denoising 과정을 역 모델에 통합함으로써 별도 스코어 계산을 없애고, 한 번의 forward‑pass만으로 깨끗한 이미지를 복원한다.

실험 결과는 ImageNet 256×256에서 DiT‑B 크기의 모델을 사용했을 때, BiFlow가 FID 2.39를 달성하면서 샘플링 속도가 최대 697배(≈2 orders of magnitude) 빨라졌음을 보여준다. 이는 기존 TARFlow 대비 4.4~2.1 FID 포인트 향상과 동시에, NF 기반 방법 중 최고 수준, 그리고 1‑NFE(단일 함수 평가) 기준에서도 경쟁력 있는 성능을 기록한다.

핵심 기여는 다음과 같다. ① 전방‑역방향을 완전히 분리해 역 모델에 가역성 제약을 없앰으로써 현대적인 비인과적 아키텍처 적용을 가능하게 함. ② hidden alignment이라는 새로운 지도 방식을 도입해 중간 trajectory 전체를 효율적으로 활용하면서도 설계 자유도를 유지. ③ score‑based denoising을 역 모델에 내재화해 인퍼런스 비용을 크게 절감. ④ 대규모 이미지 생성 벤치마크에서 기존 NF와 diffusion 기반 모델을 압도하는 품질‑속도 트레이드오프를 입증.

이러한 결과는 “학습된 역변환”이라는 오래된 아이디어가 최신 Transformer와 고성능 흐름 모델링 기법과 결합될 때, NF가 다시 현대 생성 모델의 주류로 부상할 수 있음을 시사한다. 앞으로 연속시간 흐름(Continuous Normalizing Flow)이나 Flow Matching과 같은 방법과의 시너지, 그리고 다양한 도메인(음성, 비디오, 3D)으로의 확장이 기대된다.