Beyond ReinMax: Low-Variance Gradient Estimators for Discrete Latent Variables

이 논문은 ReinMax 추정기의 높은 분산을 Rao-Blackwellisation 및 제어 변수 기법을 통해 개선한 ReinMax-Rao 와 ReinMax-CV 추정기를 제안하고, 이를 통해 이산 잠재 변수를 가진 변이 오토인코더의 학습 성능을 향상시켰음을 보여줍니다.

핵심

이 논문은 머신러닝이라는 거대한 세계를 여행하는 동안, 우리가 **'이산형 (Discrete)'**이라는 이름의 까다로운 장벽을 넘을 때 겪는 문제를 해결하는 새로운 지도를 제시합니다.

한마디로 요약하면: "정확하지만 너무 불안정한 나침반 (ReinMax) 을 찾아서, 흔들림은 줄이고 방향은 유지하는 더 나은 나침반 (ReinMax-Rao, ReinMax-CV) 을 만들었습니다."

이 내용을 일상적인 비유로 풀어보겠습니다.

---

1. 배경: 왜 우리는 '나침반'이 필요한가요?

머신러닝 모델은 보통 '연속적인' 숫자 (예: 1.5, 2.34) 로 작동합니다. 하지만 때로는 '이산적인' 선택 (예: '사과'를 고르거나 '배'를 고르는 것) 을 해야 할 때가 있습니다.

• 문제: 이산적인 선택은 '뚝' 하고 끊어지는 성질이 있어서, 수학적으로 미분 (기울기 계산) 을 할 수 없습니다. 마치 계단을 오를 때 "계단 사이의 중간 지점"을 상상할 수 없는 것과 같습니다.
• 해결책: 그래서 연구자들은 **'Straight-Through (ST)'**라는 가상의 나침반을 발명했습니다. 이 나침반은 앞에서는 실제 선택 (계단) 을 하지만, 뒤에서 학습할 때는 마치 계단이 없는 평지인 것처럼 착각하게 만들어 기울기를 계산합니다.
• 단점: 이 나침반은 편향 (Bias) 이 있어 정확한 방향을 가리키지 못하지만, 계산이 빨라 흔들림 (분산, Variance) 은 적습니다.

2. ReinMax: 정확하지만 너무 불안정한 '초고속 나침반'

최근 등장한 ReinMax라는 나침반은 이 문제를 해결하기 위해 더 정교한 수학 (Heun 의 방법, ODE) 을 사용했습니다.

• 특징: 이전 나침반보다 훨씬 **정확한 방향 (낮은 편향)**을 가리킵니다. 마치 고도계가 달린 정밀 나침반 같습니다.
• 문제: 하지만 너무 정밀하다 보니 심하게 흔들립니다 (높은 분산). 바람이 조금만 불어도 방향이 뚱뚱해져서, 모델이 학습하는 과정에서 "어디로 가야 할지" 혼란을 겪게 만듭니다.

비유: ReinMax 는 "정확한 GPS"지만 배터리가 빨리 닳고 신호가 끊겨서 지도가 자꾸 찢어지는 상황과 같습니다.

3. 이 연구의 핵심: 흔들림을 잡는 두 가지 방법

저자들은 이 "정확하지만 불안정한 나침반"을 더 안정적으로 만들기 위해 두 가지 기술을 적용했습니다.

① ReinMax-Rao: "여러 번 물어보고 평균을 내기" (Rao-Blackwellisation)

• 원리: 한 번의 측정으로 결정하는 대신, 조건에 따라 여러 번 시뮬레이션을 돌려서 평균을 내는 방식입니다.
• 비유: 길을 찾을 때 "한 명에게 물어보고 가는 것" (기존 ReinMax) 대신, "주변에 있는 10 명에게 물어보고 가장 많이 나오는 답을 고르는 것"과 같습니다.
• 결과: 흔들림 (분산) 이 확 줄어듭니다. 하지만 너무 많은 사람을 물어보는 과정에서 원래의 정확한 방향 (편향) 이 약간 흐려질 수 있습니다.

② ReinMax-CV: "보조 나침반을 붙이기" (Control Variates)

• 원리: 흔들리는 나침반 옆에, 흔들림은 적지만 방향은 약간 틀릴 수 있는 '보조 나침반'을 붙여서 서로의 오차를 상쇄시키는 기술입니다.
• 비유: 배를 항해할 때, 거친 파도에 흔들리는 큰 나침반 옆에, 작지만 안정된 나침반을 하나 더 붙여서 두 개의 방향을 평균내면 배가 덜 흔들리는 원리입니다.
• 결과: ReinMax-Rao 보다 더 균형을 잡았습니다. 흔들림은 줄이면서 원래의 정확도도 최대한 유지했습니다.

4. 실험 결과: 복잡한 미로일수록 더 효과적

연구자들은 이 새로운 나침반들을 '이산형 잠재 공간'을 가진 VAE(가상 이미지 생성기) 훈련에 적용해 보았습니다.

• 결론:
  • 단순한 문제 (작은 미로): 흔들림이 적고 편향이 있는 기존 방법들이 나쁘지 않았습니다.
  • 복잡한 문제 (거대한 미로): 차원이 높고 복잡한 상황에서는 새로운 방법 (ReinMax-Rao, ReinMax-CV) 이 압도적으로 잘 작동했습니다.
  • 이유: 복잡한 미로에서는 '정확한 방향'보다 '안정적인 진행'이 더 중요하기 때문입니다. 흔들림이 적은 나침반이 미로를 빠져나가는 데 훨씬 유리했습니다.

5. 흥미로운 발견: 왜 'Heun 의 방법'이 최고일까?

저자들은 "ReinMax 가 사용한 수학 (Heun 의 방법) 말고 다른 더 복잡한 수학 (Runge-Kutta 방법) 을 쓰면 더 정확하지 않을까?"라고 궁금해했습니다. 하지만 실험 결과는 놀라웠습니다.

• 발견: 다른 복잡한 방법들은 오히려 성능이 떨어졌습니다.
• 이유: 이 문제를 '미분 방정식 (ODE)'으로 풀려고 애쓰기보다, **'적분 (Numerical Integration)'**의 관점, 즉 **'사다리꼴 공식 (Trapezoidal Rule)'**으로 보면 이해가 쉽습니다.
  • 비유: 두 점 (시작과 끝) 을 잇는 가장 간단한 직선 (사다리꼴) 이 이미 가장 효율적인 방법이라는 뜻입니다. 더 복잡한 곡선을 그리려고 하면 오히려 계산이 복잡해지고 오차가 생깁니다.

요약 및 결론

이 논문은 머신러닝에서 이산적인 선택을 할 때 겪는 '정확성 vs 안정성'의 딜레마를 해결했습니다.

1. ReinMax는 정확하지만 너무 불안정했습니다.
2. ReinMax-Rao와 ReinMax-CV는 이 불안정함을 줄여주어, 특히 복잡하고 큰 문제를 풀 때 훨씬 뛰어난 성능을 보여주었습니다.
3. 더 복잡한 수학 공식을 쓰지 않아도, **단순하고 직관적인 수학적 도구 (사다리꼴 규칙)**가 이미 최적의 해답이었다는 것을 발견했습니다.

결국, 이 연구는 "더 복잡한 나침반을 만들려고 애쓰기보다, 흔들리는 나침반을 어떻게든 안정화시키는 것이 더 현명한 길이다"라는 교훈을 줍니다.

기술 요약

1. 연구 배경 및 문제 정의 (Problem)

• 이산 잠재 변수의 학습 난제: 머신러닝 모델에서 이산적 (discrete) 인 잠재 변수를 다룰 때, 확률 변수의 샘플링 과정은 미분 불가능 (non-differentiable) 하므로 역전파 (backpropagation) 를 통한 경사 하강법이 직접 적용되지 않습니다.
• 기존 방법의 한계:
  • Straight-Through (ST) Estimator: 계산 효율이 높고 분산 (variance) 이 낮지만, 미분 불가능한 연산을 항등 함수로 근사하는 휴리스틱 방식을 사용하여 편향 (bias) 이 큽니다.
  • ReinMax Estimator (Liu et al., 2023): 편향을 줄이기 위해 수치 ODE 관점 (Heun's method, 2 차 Runge-Kutta 방법) 에서 유도된 2 차 근사를 도입했습니다. 이는 ST 보다 편향이 적지만, 확률 변수 $D$에 의존하는 항이 포함되어 분산이 매우 높다는 치명적인 단점이 있습니다.
• 핵심 문제: ReinMax 와 같이 편향이 낮은 추정자는 분산이 너무 커서 실제 학습 (특히 VAE 등) 에서 성능이 저하될 수 있습니다. 따라서 편향을 유지하면서 분산을 줄이는 새로운 추정자가 필요합니다.

2. 제안된 방법론 (Methodology)

저자들은 ReinMax 의 높은 분산을 줄이기 위해 Rao-Blackwellisation과 Control Variates (CV) 기법을 결합하여 두 가지 새로운 추정자를 제안했습니다.

A. ReinMax-Rao (Rao-Blackwellisation 적용)

• 아이디어: ReinMax 의 식에서 높은 분산을 유발하는 첫 번째 항 ($\hat{\nabla}_{ST, \tau}(D, \theta_D)$) 을 대체합니다.
• 구현: 이 항을 Gumbel-Rao 추정자로 대체합니다. Gumbel-Rao 는 조건부 마진화 (conditional marginalisation) 를 통해 Gumbel-Softmax 재파라미터화 기법의 분산을 줄이는 것으로 알려져 있습니다.
• 효과: 분산은 크게 감소하지만, Gumbel-Rao 의 근사 특성으로 인해 ReinMax 대비 편향이 약간 증가합니다.

B. ReinMax-CV (Control Variates 적용)

• 아이디어: ReinMax-Rao 에서 발생한 편향을 보정하기 위해 Control Variates 기법을 도입합니다.
• 구현:
  • 제어 변수 (Control Variate) 로 Straight-Through Gumbel-Softmax (STGS) 추정자를 사용합니다.
  • ReinMax-Rao 식의 편향된 항을 $ST - \eta \cdot STGS + \eta \cdot E[STGS]$ 형태로 수정합니다.
  • $E[STGS]$의 닫힌 형식 해가 없으므로, 이를 Gumbel-Rao 추정자로 근사하여 계산합니다.
• 효과: 이론적으로는 편향을 보존하면서 분산을 줄여야 하나, 구현상 조건부 재파라미터화를 통한 미분을 무시하는 관행 때문에 약간의 편향이 발생하지만, ReinMax-Rao 보다는 편향이 적고 분산은 ReinMax 보다 낮아 편향 - 분산 트레이드오프의 균형점을 제공합니다.

C. 수치적 관점의 재해석 (Numerical Integration Perspective)

• 저자들은 ReinMax 를 단순히 수치 ODE (Heun's method) 의 관점이 아닌, **수치 적분 (Numerical Integration)**의 관점에서 재해석했습니다.
• 2 차 Runge-Kutta 방법의 일반화 (매개변수 $\beta$) 를 시도했으나, 실험 결과 $\beta=0.5$(Heun's method 에 해당) 일 때만 최적의 성능을 보였습니다.
• 결론: 이는 2 차 Runge-Kutta 방법보다는 **사다리꼴 법칙 (Trapezoidal Rule)**이 이 문제 (두 지점 $g'(0)$와 $g'(1)$ 사이의 적분 근사) 에 더 적합함을 의미합니다. 더 정교한 고차 다항식 근사 (예: 심슨의 법칙, 3 차 스플라인) 는 추가적인 미분 정보 (Hessian 등) 가 필요하거나 계산 비용이 너무 커서 비실용적입니다.

3. 주요 기여 (Key Contributions)

1. 새로운 추정자 제안: ReinMax 의 높은 분산 문제를 해결하기 위해 ReinMax-Rao와 ReinMax-CV를 개발했습니다.
2. 편향 - 분산 트레이드오프 최적화: 실험을 통해 ReinMax-Rao 는 분산이 가장 낮고, ReinMax-CV 는 편향과 분산 사이의 균형을 이룸을 입증했습니다.
3. 이론적 통찰: ReinMax 를 수치 ODE 가 아닌 수치 적분 (사다리꼴 법칙) 의 관점에서 이해함으로써, 왜 Heun's method 가 최적인지에 대한 직관적인 설명을 제공했습니다. 또한 더 정교한 수치 방법론이 이 문제에는 적합하지 않음을 논증했습니다.

4. 실험 결과 (Results)

• 실험 설정: MNIST 데이터셋을 사용하여 이산 잠재 변수를 가진 변이형 오토인코더 (VAE) 를 학습시켰습니다. (다양한 차원의 카테고리 및 잠재 공간 설정)
• 성능 비교:
  • 분산 (Variance): ReinMax-Rao 와 ReinMax-CV 는 원본 ReinMax 에 비해 분산을 획기적으로 감소시켰습니다 (Table 1, Figure 1).
  • 편향 (Bias): ReinMax-Rao 는 편향이 약간 증가했으나 여전히 기존 ST 기반 추정자보다 우월했습니다. ReinMax-CV 는 ReinMax-Rao 보다 편향이 낮았습니다 (Figure 1).
  • 학습 성능 (ELBO):
    • 고차원 설정 (예: 64x8, 16x12): ReinMax-Rao와 ReinMax-CV가 ReinMax 를 포함한 모든 베이스라인 (ST, Gumbel-Rao, STGS 등) 보다 우월한 ELBO 점수를 기록했습니다.
    • 저차원 설정: ReinMax(편향은 낮지만 분산이 높음) 가 단순한 문제에서는 나쁘지 않았으나, 복잡한 고차원 문제에서는 분산이 큰 추정자가 학습을 방해하는 것으로 나타났습니다.
• 결론: 저분산 추정자 (ReinMax-Rao/CV) 는 고차원 이산 잠재 공간 학습에 특히 효과적임을 입증했습니다.

5. 의의 및 결론 (Significance)

• 실용적 가치: 이산 잠재 변수를 사용하는 생성 모델 (VAE 등) 의 학습 안정성과 성능을 크게 향상시킵니다. 특히 고차원 공간에서의 학습 효율성을 높여줍니다.
• 이론적 기여: 수치 ODE 기반의 접근법이 항상 최선은 아니며, 문제의 본질에 맞는 **수치 적분 관점 (Trapezoidal Rule)**이 더 적합할 수 있음을 보여주었습니다.
• 향후 방향: 편향을 줄이면서 분산을 낮추는 더 정교한 수치적 방법론을 찾는 것은 여전히 계산 효율성 측면에서 어려운 과제로 남았으며, 본 논문은 이 방향에 대한 중요한 통찰을 제공했습니다.

요약하자면, 본 논문은 ReinMax 의 높은 분산 문제를 Rao-Blackwellisation 과 Control Variates 기법으로 해결하여, 고차원 이산 잠재 변수 모델 학습에 있어 **더 안정적이고 효율적인 새로운 표준 (ReinMax-Rao/CV)**을 제시한 연구입니다.