Learning Affine-Equivariant Proximal Operators

1. 문제 정의 (Problem Definition)

• 배경: 신호 처리 및 기계 학습, 특히 이미지 복원 (Denoising, Deblurring 등) 과 같은 역문제 (Inverse Problems) 해결에서 **근사 연산자 (Proximal Operator)**는 핵심적인 역할을 합니다. 근사 연산자는 정규화 항 (Regularizer) 을 포함하는 최적화 문제의 하위 단계로 사용됩니다.
• 기존 접근법의 한계:
  • 학습된 근사 네트워크 (LPNs): 최근 연구 (Fang et al., 2024) 는 신경망을 사용하여 데이터 기반의 정규화 함수를 학습하고, 이를 통해 정확한 근사 연산자를 계산하는 LPNs를 제안했습니다. 이는 데이터 적응성과 해석 가능성 사이의 균형을 이룹니다.
  • 구조적 제약의 부재: 그러나 기존 LPNs 는 학습된 정규화 함수에 추가적인 구조적 속성 (예: 이동 불변성, 스케일 불변성 등) 을 강제하는 방법이 명확하지 않았습니다.
  • 공변성 (Equivariance) 의 중요성: 데이터가 특정 대칭성 (예: 이미지 밝기 변화, 노이즈 레벨 변화) 을 가질 때, 이를 학습된 모델이 반영하면 일반화 성능과 강건성이 크게 향상됩니다. 특히 **스케일 (Scale)**과 **이동 (Shift)**에 대한 공변성은 노이즈 분포나 조명 조건이 훈련 데이터와 다른 환경 (Out-of-Distribution) 에서 중요합니다.
• 핵심 문제: "데이터 기반 학습이 가능하면서도, 스케일과 이동 (Affine 변환) 에 대해 공변성 (Equivariance) 을 보장하는 정확한 근사 연산자를 어떻게 설계할 수 있는가?"

2. 방법론 (Methodology)

저자들은 **Affine-Equivariant Learned Proximal Networks (AE-LPNs)**를 제안합니다. 이는 신경망으로 파라미터화된 함수가 정확한 근사 연산자이면서 동시에 Affine 공변성을 만족하도록 보장하는 프레임워크입니다.

2.1. 이론적 기반

• 근사 연산자와 볼록 함수의 관계: 어떤 함수 $f$가 볼록 함수 $\psi$의 기울기 (Gradient) 일 때 ($f = \nabla \psi$), 그 $f$는 어떤 정규화 함수 $R$에 대한 근사 연산자입니다. LPNs 는 입력이 볼록한 신경망 (ICNN) 을 사용하여 $\psi$를 학습하고, 그 기울기를 근사 연산자로 사용합니다.
• 공변성 조건:
  • 스케일 공변성: $f(ax) = af(x)$를 만족해야 합니다.
  • 이동 공변성: $f(x+b) = f(x)+b$를 만족해야 합니다.
  • Affine 공변성: $f(ax+b) = af(x)+b$를 만족해야 합니다.
• 주요 Lemma:
  1. 스케일 공변성: 함수 $\Psi$가 2 차 동차 (Homogeneous of degree 2, 즉 $\Psi(ax)=a^2\Psi(x)$) 라면, 그 기울기 $\nabla \Psi$는 1 차 동차 (스케일 공변성) 가 됩니다.
  2. Affine 공변성: 입력을 평균 성분 ($Px$) 과 평균이 제거된 성분 ($(I-P)x$) 으로 분해하여 처리하면 이동 공변성을 달성할 수 있습니다.

2.2. AE-LPN 아키텍처 설계

기존 LPN (ICNN) 을 기반으로 다음과 같은 구조적 변경을 가합니다:

1. 편향 (Bias) 제거: 모든 레이어와 잔차 연결 (Residual connection) 에서 편향 항 (Bias terms) 을 제거합니다. 이는 이동 공변성을 해치지 않기 위함입니다.
2. 활성화 함수: SortPool과 같은 스케일 공변성을 가진 비선형 함수를 사용하여, 네트워크가 1 차 동차 (Scale-equivariant) 가 되도록 합니다.
3. 2 차 동차 함수 생성:
   • 먼저 스케일 공변성을 가진 ICNN ($\Psi_\theta$) 을 학습합니다.
   • 이 네트워크의 출력을 제곱하여 $h_\theta(x) = \Psi_\theta(x)^2$로 만듭니다.
   • $\Psi_\theta$가 1 차 동차이므로, $h_\theta$는 2 차 동차가 되어, 그 기울기 $\nabla h_\theta$는 스케일 공변성을 갖게 됩니다.
4. 이동 공변성 통합:
   • 입력 $x$를 평균 벡터 $Px$와 평균 제거 성분$(I-P)x$로 분할합니다.
   • 최종 볼록 잠재 함수 $\tilde{h}(x)$를 다음과 같이 정의합니다:
     $$\tilde{h}(x) = h_\theta((I-P)x) + \frac{1}{2}\|Px\|_2^2$$
   • 이 함수의 기울기 $\nabla \tilde{h}(x)$는 Affine 공변성 ($f(ax+b) = af(x)+b$) 을 만족하며, 동시에 정확한 근사 연산자입니다.

2.3. 학습 전략

• Proximal Matching Loss (LPM): 기존 LPN 과 동일하게, 데이터의 로그 밀도 (Log-density) 를 복원하도록 학습합니다.
  • $y = x + z$ (가우시안 노이즈) 를 입력으로 하여 $f_\theta(y)$가 $x$를 복원하도록 $\ell_1$ 손실로 사전 학습 (Pre-training) 한 후, Proximal Matching Loss ($\gamma \to 0$) 로 미세 조정합니다.
• Normalization Trick 배제: 기존에 공변성을 얻기 위해 입력을 정규화 (Normalization trick) 하던 방식은 근사 연산자로서의 수학적 보장을 깨뜨리므로 사용하지 않습니다.

3. 주요 기여 (Key Contributions)

1. AE-LPN 프레임워크 제안: 신경망 파라미터화된 함수가 정확한 근사 연산자이면서 **Affine 공변성 (이동 및 스케일)**을 동시에 만족하도록 보장하는 첫 번째 방법론을 제시했습니다.
2. 이론적 증명: 2 차 동차 볼록 함수의 기울기가 스케일 공변성을 가지며, 이를 평균 분해 기법과 결합하여 이동 공변성까지 확장할 수 있음을 수학적으로 증명했습니다.
3. 기존 방법론과의 차별화:
   • 단순한 "정규화 트릭 (Normalization Trick)"은 근사 연산자 보장을 잃습니다.
   • 기존 공변성 denoiser 는 근사 연산자 (볼록 함수의 기울기) 가 아닙니다.
   • AE-LPN 은 두 가지 요구사항 (근사 연산자 + 공변성) 을 모두 충족합니다.

4. 실험 결과 (Results)

저자들은 합성 데이터와 실제 이미지 데이터 (BSDS500) 를 통해 실험을 수행했습니다.

• 합성 데이터 (Split Normal Distribution):
  • 1 차원 Split Normal 분포에서 AE-LPN 만이 진정한 근사 연산자와 로그 사전 확률 (Log-prior) 을 정확하게 학습했습니다.
  • 다른 방법들 (정규화 트릭, 스케일/이동 전용 LPN) 은 공변성이나 근사 연산자 보정 중 하나를 잃었습니다.
• 이미지 노이즈 제거 (Denoising across Noise Levels):
  • 훈련 데이터: 훈련 노이즈 레벨 (0.1) 에서는 기존 LPN 이 약간 더 높은 PSNR 을 보였습니다.
  • 테스트 데이터 (Out-of-Distribution): 훈련과 다른 노이즈 레벨 (0.05 ~ 0.7) 에서 **AE-LPN 이 가장 높은 강건성 (Robustness)**을 보였습니다. 특히 스케일 공변성이 노이즈 레벨 변화에 더 중요한 영향을 미쳤습니다.
• Affine 변환 검증:
  • 이미지 밝기 변화 ($g(x) = \alpha x + (1-\alpha)$) 를 적용했을 때, AE-LPN 은 입력의 변환에 따라 출력이 정확히 변환되는 공변성을 완벽하게 유지했습니다 (PSNR 이 매우 높음).
  • 반면, 다른 모델들은 변환에 대한 불일치를 보였습니다.

5. 의의 및 결론 (Significance & Conclusion)

• 강건한 역문제 해결: AE-LPN 은 훈련 분포를 벗어난 환경 (다른 노이즈 레벨, 조명 변화 등) 에서도 안정적인 성능을 발휘하여, 실제 응용 (의료 영상, 위성 영상 등) 에서의 실용성을 크게 높입니다.
• 수학적 엄밀성: "학습된 모델"이 단순히 흑상자 (Black-box) 가 아니라, 수학적 정의 (볼록성, 근사 연산자 성질) 를 엄격하게 준수하면서도 데이터의 구조적 속성 (공변성) 을 학습할 수 있음을 입증했습니다.
• 미래 전망: 이 프레임워크는 단순한 노이즈 제거를 넘어, CT 재구성 (Windowing 및 Intensity scaling 변화 대응) 과 같은 다른 역문제 및 샘플 복잡도 감소에도 확장 가능할 것으로 기대됩니다.

요약하자면, 이 논문은 수학적 보장을 유지하면서 데이터의 기하학적 구조 (Affine 변환) 를 학습에 효과적으로 통합하는 새로운 신경망 아키텍처를 제안하여, 역문제 해결의 강건성과 일반화 능력을 획기적으로 향상시켰습니다.