기술 요약: FlowADMM: 플로우 기반 재노이즈-디노이즈 사전 모델을 활용한 플러그앤플레이 ADMM

문제 정의

이미지 복원, 디블러링, 초해상도, 인페인팅과 같은 이미징의 역문제들은 일반적으로 잘 정의되지 않은(ill-posed) 문제이며, 열화된 측정값 $y = Ax + \nu$로부터 이미지 $x$를 복원하는 것으로 공식화됩니다. 딥러닝은 강력한 직접 매핑을 제공하지만, 종종 해석 가능성과 이론적 보장이 부족합니다. 고전적인 변분 접근법은 이를 데이터 충실도 항과 정규화 항의 합을 최소화하는 문제로 공식화하며, 이는 Forward-Backward Splitting (FBS) 과 같은 근사 알고리즘을 통해 해결됩니다.

Plug-and-Play (PnP) 방법은 수동으로 설계된 정규화 함수의 근사 연산자를 학습된 디노이저로 대체함으로써 이를 개선합니다. 최근 생성 플로우 (generative flow) 및 확산 모델은 특정 노이즈 수준에 적응된 시간 의존적 디노이저 계열 $\{D_t\}$를 제공했습니다. 그러나 주요 과제가 발생합니다: 반복적 재구성 과정에서 중간 반복값은 디노이저가 훈련된 노이즈 분포 $p_t$에서 종종 벗어납니다. 기존 플로우 기반 PnP 방법 (예: PnP-Flow, Flower) 은 명시적인 재노이즈 - 디노이즈 (renoise-denoise) 단계를 통합하여 이를 해결합니다. 이는 디노이징 전에 현재 반복값을 적절한 노이즈 수준으로 교란시키는 과정입니다. 이러한 방법들은 종종 과정을 안정화시키기 위해 여러 노이즈 샘플에 대해 평균을 냅니다.

경험적 성공에도 불구하고, 기존 재노이즈 - 디노이즈 PnP 방법들은 확률적 알고리즘으로 공식화되어 이론적으로 수렴 거동을 분석하기 어렵습니다. 구체적으로, 잠재 노이즈 분포에 걸쳐 디노이저 출력의 평균을 취함으로써 암묵적으로 유도된 결정론적 연산자는 아직 특징화되지 않았으며, 이로 인해 이러한 방법들의 수렴 특성은 잘 이해되지 않고 있습니다.

방법론

평균 재노이즈 - 디노이즈 연산자

저자들은 확률적 플로우 기반 PnP 방법의 기반이 되는 결정론적 연산자를 식별하고 공식화했습니다. 시간 의존적 MMSE 디노이저 $D_t = I + (1-t)v_\theta^t$를 정의하는 사전 훈련된 플로우 속도 네트워크 $v_\theta$가 주어졌을 때, 표준 확률적 단계는 노이즈 $\epsilon \sim \mathcal{N}(0, I)$를 사용하여 반복값 $x$를 재노이즈하여 $tx + (1-t)\epsilon$를 얻은 후 디노이징하는 것을 포함합니다.

본 논문은 이 과정의 노이즈 분포에 대한 기댓값으로 평균 재노이즈 - 디노이즈 연산자 $\bar{S}_t(x)$를 정의합니다:
$$\bar{S}_t(x) = \mathbb{E}_{\epsilon \sim \mathcal{N}(0, I)} [D_t(tx + (1-t)\epsilon)]$$
이 연산자는 PnP-Flow 와 같은 방법에서 사용되는 확률적 업데이트의 결정론적 극한을 나타냅니다. 저자들은 잔차 형태 $R_t(x) = \bar{S}_t(x) - x$를 분석하여, 이를 데이터 매니폴드를 향한 국소 보정을 제공하는 벡터장으로 해석합니다. 이상적인 MMSE 가정 하에서, 데이터 매니폴드 위의 점들은 이 연산자의 고정점입니다.

FlowADMM 알고리즘

이 통찰력을 바탕으로 저자들은 FlowADMM을 제안했는데, 이는 평균 재노이즈 - 디노이즈 연산자를 승수 교대 방향법 (ADMM) 프레임워크에 통합합니다. FBS 와 달리, ADMM 은 여러 사전 연산자를 허용하며 연산자 의존적 스텝 크기 제한 없이 비미분 가능한 데이터 항을 처리합니다.

이 알고리즘은 제약 조건 $x=z$ 하에서 $\min_{x,z} F_y(x) + R(z)$ 문제를 확장된 라그랑지안을 최소화함으로써 해결합니다. 업데이트는 다음과 같습니다:

1. 데이터 일관성: $x^{k+1} = \text{prox}_{\tau F_y}(z^k - u^k)$
2. 사전 업데이트: $z^{k+1} = \bar{S}_{t_k}(x^{k+1} + u^k)$
3. 이중 업데이트: $u^{k+1} = u^k + x^{k+1} - z^{k+1}$

$\bar{S}_t$의 기댓값은 계산 불가능하므로 몬테카를로 샘플링을 통해 근사됩니다. 저자들은 시간 단계 의존적 샘플링 전략을 제안합니다: 초기 반복에서는 (거친 보정이 이루어지고 연산자가 노이즈에 덜 민감할 때) 샘플 수를 줄이고, 후기 반복에서는 (세부 사항이 정제되고 정밀도가 중요할 때) 샘플 수를 늘리는 방식입니다.

수렴 분석

본 논문은 기반 플로우 네트워크의 약한 리프시츠 조건 하에서 FlowADMM 에 대한 수렴 보장을 확립합니다:

• 고정 시간 스케줄: 일정한 시간 단계 $t$에 대해, 플로우 네트워크의 리프시츠 상수 $L_v$가 $L_v < \frac{1}{1-t}$를 만족하면 수렴이 보장됩니다. 이 조건은 많은 다른 PnP 방법에서 요구되는 엄격한 축소 ($L < 1$) 조건보다 훨씬 약하며, 특히 $t \to 1$일 때 그렇습니다.
• 비정상 시간 스케줄: 분석은 $t_{max} < 1$로 수렴하는 시간 스케줄 $(t_k)_k$로 확장됩니다. 저자들은 연산자 열이 충분히 빠르게 극한에 접근하고 플로우 네트워크가 시간 $t$에 대해 리프시츠 연속일 때, 극한 연산자 $T_{t_{max}}$의 고정점으로 수렴함을 증명합니다.

주요 기여

1. 이론적 특징화: 저자들은 확률적 플로우 기반 PnP 방법의 기반이 되는 결정론적 평균 재노이즈 - 디노이즈 연산자를 식별하고 분석하여, 그 잔차 구조와 리프시츠 거동을 특징화했습니다.
2. 알고리즘 제안: 그들은 이 평균 연산자를 ADMM 프레임워크에 직접 통합하는 PnP 알고리즘인 FlowADMM 을 제안했습니다.
3. 수렴 보장: 그들은 약한 리프시츠 가정 하에서 FlowADMM 에 대한 수렴 기준을 유도했으며, 이는 일정하고 일정하지 않은 (시간 가변) 시간 단계 스케줄 모두를 포괄합니다.
4. 샘플링 전략: 그들은 평균 연산자의 거동에 영감을 받아 후기 반복에서 샘플 수를 늘리는 거칠기에서 정밀함으로의 (coarse-to-fine) 몬테카를로 샘플링 전략을 제안했습니다.
5. 경험적 성능: 그들은 이전 확률적 평균 접근법보다 적은 데이터 일관성 평가로 여러 역문제에서 최첨단 성능을 입증했습니다.

실험 결과

저자들은 CelebA 및 AFHQ-Cat 데이터셋에서 디노이징, 디블러링, 초해상도, 무작위 인페인팅, 박스 인페인팅 등 다섯 가지 작업에 대해 FlowADMM 을 평가했습니다.

• 성능: FlowADMM 은 최근 플로우 및 확산 기반 PnP 방법 (PnP-Flow, Flower, OT-ODE, D-Flow, DiffPIR 포함) 과 비교하여 10 개 작업 중 7 개에서 왜곡 지표 (PSNR 및 SSIM) 에서 최첨단 결과를 달성했습니다. 초해상도 및 인페인팅 작업에서 주목할 만한 개선이 관찰되었습니다.
• 효율성: 주된 경쟁자 인 Flower5-OT 의 총 플로우 평가 (네트워크 순전파) 수와 일치하지만, FlowADMM 은 데이터 항 근사 연산자 평가 횟수가 훨씬 적었습니다. 이는 Flower5-OT 가 전체 재구성 궤적을 평균내는 반면, FlowADMM 은 각 반복 내에서 사전 단계만 평균내기 때문입니다.
• 애블레이션 연구: 실험은 "후기 중점 (late-heavy)" 샘플링 스케줄 (후기 반복에서 샘플 증가) 이 일정한 샘플링보다 성능을 크게 향상시킨다는 것을 확인했습니다. 또한 이 스케줄을 PnP-Flow 에 적용하면 결과가 개선되었지만, FlowADMM 이 여전히 더 우수하여 ADMM 공식화 자체가 이점을 제공함을 시사했습니다.
• 한계: 수렴 분석은 훈련 중에 명시적으로 패널티가 부과되지 않는 사전 훈련된 네트워크의 리프시츠 상수에 의존합니다. 또한, 이론적 보장은 결정론적 평균 연산자에 적용되지만, 실제 구현은 확률적 몬테카를로 근사를 사용합니다.

중요성 및 주장

본 논문은 확률적 재노이즈 - 디노이즈 방법의 기반이 되는 결정론적 연산자를 명시적으로 공식화함으로써, 플로우 기반 PnP 에서 경험적 성공과 이론적 이해 사이의 간극을 해소했다고 주장합니다. 이 연구는 다음을 제공합니다:

• 검증 가능한 수렴: 이전 확률적 PnP 방법과 달리, FlowADMM 은 플로우 네트워크에 대한 현실적인 리프시츠 조건 하에서 수렴 보장을 제공합니다.
• 알고리즘적 효율성: 궤적 평균화 방법과 비교하여 ADMM 에의 통합은 계산 자원의 더 효율적인 사용 (적은 데이터 일관성 단계) 을 가능하게 합니다.
• 연산자 이론적 통찰: 분석은 평균 연산자의 "거칠기에서 정밀함으로" 거동을 적응형 샘플링 스케줄을 통해 활용할 수 있음을 시사하며, 이 원칙은 ADMM 과 FBS 기반 접근법 모두에 이점을 줍니다.

저자들은 FlowADMM 이 플로우 기반 사전 모델을 가진 역문제를 해결하기 위한 견고하고 이론적으로 근거 있는 프레임워크를 제공한다고 결론지었으며, 향후 연구는 훈련 중 리프시츠 제약을 부과하여 안정성 보장을 강화하는 데 초점을 맞출 수 있음을 인정했습니다.