Amortized Factor Inference Networks for Posterior Inference

기술적 요약: 상각 인자 추론 네트워크 (AFINs)

1. 문제 제기

상각 추론 (Amortized inference) 은 관측치를 사후 분포로 매핑하는 단일 신경망을 학습하여, 일회성 훈련 비용 이후 빠른 테스트 시간 추론을 제공하는 것을 목표로 합니다. 그러나 기존 상각 방법들은 일반적으로 고정된 확률 모델에 종속되어 있습니다. 만약 사전 분포 (prior), 가능도 (likelihood), 잠재 차원 ($d$), 또는 관측치 수 ($N$) 가 변경되면, 추론 네트워크는 보통 재훈련이나 비용이 큰 테스트 시간 미세 조정이 필요합니다. 일부 최근 연구들은 변화하는 $N$이나 특정 매개변수 변경에 대해 일반화하지만, 정확도를 희생하거나 모델별 적응을 요구하지 않으면서 동시에 변화하는 사전/가능도 계열, 차원, 그리고 관측치 수의 임의의 조합을 처리하는 능력은 종종 부족합니다.

이 연구가 제기하는 핵심 질문은 다음과 같습니다: 단일 추론 네트워크가 다양한 사전 분포, 가능도, 차원, 그리고 관측치 수를 처리하면서도 NUTS 와 같은 비상각 방법의 정확도와 일치하는 상태로, 전체적인 다양한 베이지안 추론 문제 군에서 일반화할 수 있을까요?

2. 방법론: 상각 인자 추론 네트워크 (AFINs)

저자들은 모델 명세와 그 관측치를 직접 변분 사후 분포의 매개변수로 매핑하도록 설계된 인코딩 - 병합 - 디코딩 아키텍처인 **상각 인자 추론 네트워크 (AFINs)**를 소개합니다.

A. 타입 지정 인자 표현 (Typed-Factor Representation)
AFIN 은 모델을 단일 블록으로 취급하는 대신, 결합 밀도 $p(z, y_{1:N})$을 $N+1$개의 인자 집합으로 표현합니다: 하나의 사전 인자 $p_0(z)$와 $N$개의 가능도 인자 $p_n(y_n|z)$입니다. 각 인자는 다음과 같이 표현됩니다:

• 타입 ($t_n$): 분포 계열을 나타내는 레이블 (예: "가우시안 가능도", "스튜던트-t 가능도").
• 매개변수 ($\theta_n$): 인자를 인스턴스화하는 데 필요한 인자별 매개변수 (예: 스케일, 자유도).
• 관측치 ($y_n$): 가능도 인자와 연관된 데이터.

이 표현은 서로 다른 관측치가 서로 다른 가능도 계열을 따를 수 있는 이질적 모델을 처리할 수 있게 합니다.

B. 차원 독립 아키텍처
변화하는 잠재 차원 ($d$) 과 인자 수 ($N$) 를 지원하기 위해, AFIN 은 $d$에 독립적이고 잠재 좌표의 치열에 대해 공변 (equivariant) 인 모듈을 사용합니다:

1. BoxMLP: 학습된 가중치가 $d$에 의존하지 않는 매개변수 결속 (parameter-tied) 다층 퍼셉트론입니다. 좌표별 특징과 풀링된 요약 (예: 차원 평균) 을 처리하여 출력물이 잠재 벡터 $z \in \mathbb{R}^d$의 치열에 대해 공변하도록 보장합니다.
2. BoxTransformer: 표준 Transformer 의 차원 독립 아날로그입니다. 인자 임베딩 시퀀스에 작동합니다. 결정적으로, 프로젝션 맵 (Query, Key, Value, Output, FFN) 은 좌표 인덱싱 토큰에 적용되는 BoxMLP 로 구현됩니다. 이는 잠재 차원을 평탄화하지 않고 인자 간 상호작용을 모델링할 수 있게 하여, 매개변수의 $O(d)$ 또는 $O(d^2)$ 구조를 보존합니다.

C. 인코딩 - 병합 - 디코딩 파이프라인

1. 인코딩: 각 인자는 독립적으로 처리됩니다. 타입별 어댑터가 인자의 매개변수와 관측치를 공통 임베딩 공간으로 매핑한 후, 공유 BoxMLP 인코더가 이를 처리합니다. 임베딩에는 노드 특징 (좌표별) 과 쌍 특징 (좌표 쌍 상호작용, 공분산/정밀도 행렬에 유용) 이 모두 포함됩니다.
2. 병합: 인자 임베딩 시퀀스는 $M$개의 BoxTransformer 블록에 의해 처리됩니다. 이러한 블록은 인자 축을 가로지르는 어텐션 (인자를 시퀀스 차원으로 취급) 을 사용하여 정보를 집계합니다. 어텐션 메커니즘은 노드 점수와 쌍 점수를 결합하여 표현을 업데이트합니다.
3. 디코딩: 병합된 표현은 인자 축을 따라 합계 풀링되어 단일 작업 수준 임베딩을 생성합니다. 최종 BoxMLP 디코더는 이 임베딩을 변분 매개변수 $\phi$ (예: 가우시안의 평균과 정밀도, 또는 정규화 흐름의 흐름 매개변수) 로 매핑합니다.

D. 훈련 목적 함수
AFIN 은 시뮬레이션 기반 추론을 통해 훈련됩니다. 네트워크는 생성된 관측치 $y$에 조건 지어진 결합 모델 $p(z, y)$에서 샘플링된 잠재 표본 $z$에 할당된 로그 밀도를 최대화하도록 훈련됩니다. 이는 실제 사후 분포와 변분 근사 사이의 기대 전방향 KL 발산 ($KL(p(z|y) || q_\phi(z))$) 을 최소화하는 것과 동일합니다. 전방향 KL 목적 함수는 중요도 샘플링을 위한 제안 분포로 근사를 사용할 때 유익한 것으로 지적됩니다.

3. 주요 기여

• 일반화: 단일 훈련된 AFIN 은 재훈련 없이 다양한 사전 계열, 가능도 계열, 잠재 차원 ($d$), 그리고 관측치 수 ($N$) 에 걸쳐 사후 추론을 수행할 수 있습니다.
• 차원 독립성: BoxMLP 와 BoxTransformer 모듈의 도입으로 네트워크는 고정된 매개변수 수를 유지하면서 임의의 차원 $d$의 잠재 공간에서 작동할 수 있게 되었습니다. 이는 고정된 입력/출력 크기를 요구하는 기존 상각 네트워크와 큰 차이점입니다.
• 인자 수준 추상화: 고정된 모델 그래프 대신 타입 지정 인자에 작동함으로써, AFIN 은 단일 추론 실행 내에서 이질적 가능도 (예: 가우시안과 베르누이 관측치 혼합) 를 지원합니다.
• 효율성: 이 방법은 반복적 방법들에 비해 테스트 시간 계산량을 현저히 줄이면서도 높은 정확도를 달성합니다.

4. 실험 결과

저자들은 미세 조정 없이 합성 및 실제 세계 작업에 대해 단일 5 백만 매개변수 AFIN 체크포인트를 평가했습니다.

• 합성 성능: 48 가지 설정 (16 가지 사전 - 가능도 조합 $\times$ 3 가지 난이도 수준) 전반에 걸쳐, AFIN 은 NUTS 와 다양한 변분 추론 (VI) 기준 (Full-Rank Gaussian, IAF, MAF, NSF) 과 비교 가능한 사후 정확도를 달성했습니다.
  • 속도: AFIN 단일 샷 추론은 NUTS 및 흐름 기반 VI 방법에 비해 2~4 차수 (orders of magnitude) 더 빠릅니다.
  • 정제: 자기 정규화 중요도 샘플링 (SNIS) 을 위한 제안 분포로 사용될 때, AFIN+SNIS 는 실질적으로 더 짧은 월클락 시간으로 NUTS 정확도를 일치하거나 초과했습니다.
• 외삽: AFIN 은 훈련 중 관찰된 것 ($d \le 16, N \le 256$) 보다 큰 잠재 차원 ($d=32$) 과 관측치 수 ($N=512$) 를 포함하는 분포 외 (OOD) 설정에 대한 제로샷 일반화를 보여주었습니다.
• 실제 세계 전이: 12 개의 UCI 데이터셋 (이질적 가능도 설정 포함) 에서 AFIN+SNIS 는 장기간 NUTS 참조와 거의 즉시 경쟁력 있는 정확도에 도달한 반면, 반복적 기준들은 수렴하기 위해 상당한 최적화 시간이 필요했습니다.
• 디코더 변형: 가우시안 및 흐름 디코더 모두 테스트되었습니다. 흐름 디코더는 상각의 속도 이점을 유지하면서 가우시안 디코더보다 약간 더 높은 정확도 개선을 제공했습니다.

5. 중요성 및 주장

이 논문은 AFIN 이 상각 추론의 속도와 비상각 방법의 유연성 사이의 간극을 성공적으로 연결한다고 주장합니다. 타입 지정 인자 인터페이스와 차원 독립 모듈을 통해 추론 네트워크를 고정된 모델 명세에서 분리함으로써, AFIN 은 재사용 가능한 상각 추론을 가능하게 합니다.

저자들은 AFIN 을 다양한 확률적 프로그램을 처리할 수 있는 베이지안 추론을 위한 "기초 모델"로 나아가는 한 걸음으로 위치시킵니다. 그들은 사전 정의된 인자 타입 집합의 필요성과 쌍 구성 요소 임베딩의 $O(d^2)$ 확장성과 같은 한계를 인정하지만, 이 접근 방식이 이전에는 별개의 재훈련된 모델이 필요했던 광범위한 추론 문제 군을 단일 네트워크로 해결할 수 있게 한다는 점을 강조합니다. 이 연구는 상각이 최첨단 샘플링 방법 (NUTS) 의 품질을 달성하면서도 테스트 시간에 막대한 계산적 절감을 제공할 수 있음을 시사합니다.