기술적 요약: f-궤적 균형 (f-Trajectory Balance)

문제 제기

대규모 언어 모델 (LLM) 과 생성 흐름 네트워크 (GFlowNets) 를 포함한 생성 모델의 정렬 (alignment) 및 미세 조정 (fine-tuning) 은 일반적으로 모델의 분포와 목표 분포 간의 불일치를 최소화하는 데 의존합니다. 강화 학습 (RL) 미세 조정에서는 이를 종종 Kullback-Leibler (KL) 발산 페널티를 받는 기대 보상 최대화 문제로 표현합니다.

이 분야의 주요 과제는 **모드 추구 (mode-seeking)**와 모드 커버링 (mode-covering) 행동 간의 트레이드오프입니다. 역 KL 발산 (보통 궤적 균형 또는 제곱 KL 손실을 통해 구현됨) 과 같은 표준 목적 함수는 모드 추구 성향을 띠어, 모델이 소수의 고확률 모드에 수렴하게 만듭니다. 이는 신약 개발이나 다양한 후보 생성과 같은 탐색이 필요한 작업에는 부적합하며, 이러한 작업에서는 목표 분포의 전체 다양성을 포착하는 "모드 커버링"이 선호됩니다.

또한, 기존 방법들은 **오프-폴리시 유효성 (off-policy validity)**에 종종 어려움을 겪습니다. REINFORCE 나 PPO 와 같은 추정기는 일반적이지만 높은 분산을 겪습니다. 최근 제곱 KL 손실 ($KL_{sq}$) 을 사용하는 접근법은 온-폴리시에서 낮은 분산의 기울기를 제공하고 오프-폴리시에서도 유효하지만, 본질적으로 역 KL 발산의 특성에 묶여 있어 탐색 - 활용 트레이드오프를 제어하는 능력이 제한적입니다.

방법론

저자들은 **f-궤적 균형 (f-Trajectory Balance)**을 제안합니다. 이는 제곱 KL 손실을 f-발산 (f-divergences) 전체 계열로 일반화한 일련의 대리 손실 함수 (surrogate loss functions) 입니다.

이론적 기반

핵심 이론적 기여는 로그 확률에 대한 **변환 불변 손실 함수 (translation-invariant loss functions)**와 f-발산 사이의 일대일 대응을 확립하는 것입니다.

1. 발산에서 손실로: 임의의 볼록 함수 $f$가 정의하는 f-발산 $D_f(p_\theta \| p^*)$에 대해, 저자들은 로그 확률의 차이 $\Delta_\theta(y) = \log p_\theta(y) - \log p^*(y)$에 정의된 손실 함수 $L_f$를 유도합니다:
   $$L_f(\Delta_\theta(y)) = \int_0^{\Delta_\theta(y)} (f'(e^t) - f'(1)) dt$$
   이 손실은 두 가지 중요한 특성을 가집니다:

   • 온-폴리시 기울기 동등성: 샘플이 모델 정책 ($\mu = p_\theta$) 에서 추출될 때, $L_f$의 기대 자동 미분 기울기는 해당 f-발산 $D_f$의 기울기와 일치합니다.
   • 오프-폴리시 유효성: 행동 분포 $\mu$에서 샘플링된 오프-폴리시 데이터에서 평가될 때, 목표의 지원 집합이 행동 정책의 지원 집합에 포함된다면, 손실은 동일한 전역 최소화자 ($p_\theta = p^*$) 를 갖는 유효한 목적 함수로 남습니다.

2. 손실에서 발산으로: 반대로, 로그 확률에 대한 임의의 엄격하게 볼록하고 변환 불변인 손실 $\ell$은 특정 f-발산을 최소화하는 것과 대응됩니다. 이 논문은 이러한 매핑이 서로의 역함수임을 증명합니다.

실제 구현

훈련 시 발생하는 실제적인 과제를 해결하기 위해 저자들은 몇 가지 확장을 도입합니다:

• DevGrad 손실 (배치별 정규화): 분할 함수 (GFlowNets 와 비정규화 목표에서 흔히 발생) 의 처리 불가능성을 해결하기 위해, 저자들은 VarGrad 추정기를 일반화합니다. 배치 분산을 사용하는 대신, 로그 확률 차이의 배치별 일반화된 편차를 사용합니다. 이는 배치 전체에서 손실을 최소화하여 정규화 상수 $\log Z$를 추정함으로써 점수 함수 계수를 중심화하고 분산을 줄이는 방식입니다.
• 온도 조절 손실 (Tempered Loss): 작은 온도 매개변수 $\beta$ (목표 분포가 $p^* \propto \pi_{ref} \exp(r/\beta)$인 경우) 를 다룰 때 발생하는 수치적 오버플로우를 방지하기 위해, 저자들은 온도 조절 손실을 제안합니다. 이는 손실을 $1/\beta$로 스케일링하고 온도 조절된 분포에 적용하여 기울기가 유계되고 안정적으로 유지되도록 합니다.
• GFlowNets 를 위한 f-궤적 균형: 이 프레임워크는 $L_f$ 손실을 사용하여 궤적 균형 제약을 강제함으로써 GFlowNets 에 적용됩니다. 이는 표준 궤적 균형 (역 KL 에 해당) 을 임의의 f-발산으로 일반화합니다.

주요 기여

1. 일반화된 손실 계열: 해당 f-발산의 온-폴리시 기울기와 일치하면서 오프-폴리시 유효성을 유지하는 변환 불변 대리 손실 ($L_f$) 의 일반 계열 유도.
2. 이론적 동등성: 변환 불변 로그 확률 손실과 f-발산 사이의 역관계 증명. 이를 통해 원하는 특성 (예: 모드 커버링) 을 가진 새로운 손실 함수를 체계적으로 설계 가능.
3. 비정규화 목표로의 확장: DevGrad 손실을 통해 VarGrad 기법을 임의의 f-발산으로 일반화하여, RL 미세 조정에서의 처리 불가능한 분할 함수 문제를 해결.
4. 탐색 제어: f-발산의 선택 ( $\alpha$-발산에서 $\alpha$ 매개변수로 파라미터화됨) 이 모드 추구 대 모드 커버링 행동을 직접적으로 제어함을 입증. 낮은 $\alpha$ 값 (예: Forward KL, Hellinger) 은 모드 커버링을 촉진하고, 높은 값 (예: Pearson $\chi^2$) 은 모드 추구를 촉진합니다.

실험 결과

저자들은 네 가지 다른 작업에 걸쳐 손실 계열을 평가했습니다:

1. 합성 그리드 (GFlowNets): 모드 커버링을 테스트하도록 설계된 2D 그리드 작업에서, 낮은 $\alpha$ (Forward KL, Hellinger) 를 가진 손실들은 표준 궤적 균형 (역 KL) 보다 훨씬 빠르고 신뢰성 있게 목표 분포의 네 가지 모드를 모두 발견했습니다. 반면, Pearson $\chi^2$ ($\alpha=2$) 은 강한 모드 추구 행동을 보이며 처음 발견된 모드에 갇혔습니다.
2. 분자 발견 (SynFlowNet): 합성적으로 접근 가능한 분자를 생성하는 SynFlowNet 프레임워크에 적용되었습니다. 저자들은 $\alpha$를 조정함으로써 생성된 분자의 다양성을 제어할 수 있음을 발견했습니다. (탐색을 위해 낮은 $\alpha$로 시작하여 활용을 위해 $\alpha=1$로 증가시키는) 어닐링 스케줄이 가장 좋은 트레이드오프를 보여주어, 표준 궤적 균형에 비해 더 높은 보상을 주는 분자를 더 큰 다양성으로 생성했습니다.
3. 확산 모델: MNIST 숫자 (홀수 대 짝수) 의 조건부 샘플링에서, 표준 궤적 균형은 목표 모드 전체에 고르게 샘플링하지 못했습니다 (특정 숫자 0 과 6 과 같은 것을 과다 샘플링). f-궤적 균형 손실, 특히 낮은 $\alpha$를 가진 손실들은 목표 모드 전반에 걸쳐 더 균일한 샘플링을 달성했습니다.
4. 비동기 LLM 미세 조정: 이 손실들은 수학 문제 (GSM8k, MATH) 에 대해 50 단계 비동기 지연으로 LLM(Qwen 및 OLMo 계열) 을 미세 조정하는 데 적용되었습니다. f-발산 손실들은 안정적인 훈련과 일관된 엔트로피 - 보상 트레이드오프를 보인 반면, 표준 PPO 구현은 동일한 비동기 조건 하에서 불안정성을 나타냈습니다.

중요성 및 주장

이 논문은 제곱 KL 손실의 효과성이 고립된 현상이 아니라 더 넓은 손실 계열의 특정 사례임을 주장합니다. "온-폴리시 기울기 일치"와 "오프-폴리시 유효성"의 특성을 f-발산 전체 계열로 확장함으로써, 저자들은 생성 모델을 조정하기 위한 통합 프레임워크를 제공합니다.

주요 중요성은 기본 훈련 알고리즘을 변경하지 않고도 다른 대리 손실 함수를 선택하거나 (또는 $\alpha$ 매개변수를 조정하여) 탐색 - 활용 트레이드오프를 명시적으로 제어할 수 있다는 점에 있습니다. 이는 실무자들이 보상 최대화를 위한 모드 추구와 다양성 및 탐색을 위한 모드 커버링 행동 사이에서 선택할 수 있게 하며, 동시에 LLM 과 GFlowNets 의 현대적이고 대규모인 비동기 훈련에 필요한 안정성과 오프-폴리시 호환성을 유지합니다. 이 작업은 고보상 후보의 공간을 커버하는 것이 중요한 신약 개발 및 탐색 에이전트와 같은 응용 분야에 특히 가치 있음을 시사합니다.