기술 요약: 최적화된 텍스트 생성을 위한 세미 오프라인 강화 학습 (Semi-Offline Reinforcement Learning)

문제 정의

텍스트 생성을 위한 강화 학습(RL)은 현재 **온라인(online)**과 오프라인(offline) 설정 사이의 이분법적 상황에 직면해 있으며, 각 설정은 상당한 트레이드오프를 가집니다:

• 온라인 RL: 언어 모델(LM)이 자신의 출력 분포로부터 샘플링함으로써 환경을 완전히 탐색할 수 있게 합니다. 이는 최종 목적 함수(예: ROUGE, 인간 피드백)에 대한 포괄적인 이해를 제공하지만, 과도한 최적화 비용을 발생시킵니다. 구체적으로, $K$개의 샘플을 최적화하려면 모델에 대해 $K \times T$번의 순전파(forward propagation, FP)가 필요하며, 여기서 $T$는 시퀀스 길이입니다. 이는 대규모 사전 학습된 모델의 경우 실용적이지 못합니다.
• 오프라인 RL: 정적 데이터셋(예: 정답 또는 빔 서치 출력)을 활용하여 훈련 중 자기회귀적(autoregressive) 생성을 제거합니다. 이는 비용을 $KT$ FP에서 $K$ FP로 줄여줍니다. 그러나 탐색 능력은 희생됩니다. 모델은 정적 데이터 포인트에 대해서만 보상 신호를 받으므로, 정적 분포 외부의 최적의 생성 방향을 발견할 수 없습니다.

본 연구가 다루는 근본적인 연구 질문은 다음과 같습니다: 최소한의 최적화 비용으로 효과적인 탐색을 달able히 하는 정교한 RL 설정을 구축할 수 있는가?

방법론: 세미 오프라인 RL (Semi-Offline RL)

저자들은 하이퍼파라미터 $p_m \in [0, 1]$에 의해 제어되는 연속적인 설계 공간을 정의함으로써 온라인과 오프라인 설정 사이의 간극을 메우는 새로운 패러다임인 세미 오프라인 RL을 제안합니다.

1. 세미 오프라인 MDP 정식화

이 방법은 언어 모델이 생성한 토큰과 정적 데이터셋의 토큰을 혼합하여 샘플을 구성하는 마르코프 결정 과정(MDP)에 기반합니다.

• 상태 ($s_t$): 입력 $x$, 부분적으로 도출된 출력, 그리고 바이너리 마스크 시퀀스 $M_0, \dots, M_t$를 포함합니다.
• 전이(Transition): 각 단계 $t$에서 베르누이 변수 $M_{t+1} \sim \text{Bernoulli}(p_m)$가 다음 토큰의 출처를 결정합니다:
  • 만약 $M_{t+1} = 1$이면: 에이전트는 LM으로부터 토큰 $\hat{y}_t$를 샘플링합니다 (탐색).
  • 만약 $M_{t+1} = 0$이면: 에이전트는 정적 데이터셋에서 토큰 $y_t$를 가져옵니다 (활용).
• 비용 분석: 완전 관측 가능한 시나리오에서 기대 FP 횟수는 $E(C_t) = t \cdot p_m$입니다. $p_m \in (0, 1)$로 설정함으로써, 이 방법은 탐색 능력과 중간 단계의 최적화 비용 사이의 균형을 맞춥니다.

2. 최적의 RL 설정: 마스크된 관측 (Masked Observations)

최소한의 최적화 비용(구체적으로 샘플 수와 상관없이 1 FP)을 유지하면서 탐색을 수행하기 위해, 저자들은 부분 관측 MDP (POMDP) 정식화로 전환합니다.

• 제약 조건: 병렬 생성(1 FP)을 가능하게 하기 위해, 에이전트의 행동은 모델에 의해 동시에 생성되는 토큰에 의존할 수 없습니다.
• 최적의 관측 ($o^*_t$): 저자들은 최적의 관측이 세 가지 바람직한 속성(DP1-DP3)을 만족함을 증명합니다:
  1. 최소 최적화 비용: 1 FP를 통해 달성됩니다.
  2. 최소 점근적 편향 (Minimum Asymptotic Bias): 상태를 예측하는 데 필요한 모든 정보를 유지해야 합니다.
  3. 최소 과적합 오차 경계 (Minimum Overfitting Error Bound): 관측값 내의 불필요한 정보를 최소화해야 합니다.
• 구현: 최적의 설정은 **마스크된 관측(Masked Observations)**에 의해 실현됩니다. LM의 입력은 소스 $x$와, 모델이 생성한 토큰은 특수 마스크 토큰 $[M]$으로 대체되고 정적 데이터셋의 토큰은 그대로 보이는 타겟 시퀀스로 구성됩니다.
  • $y^M_t = [M]$ (모델이 생성할 때, $M_t=1$)
  • $y^M_t = y_t$ (정적 데이터 사용 시, $M_t=0$)
  • 이를 통해 모델은 관측값이 샘플 간에 동일하게 유지되므로, 단 한 번의 순전파로 $K$개의 서로 다른 궤적을 동시에 탐색할 수 있습니다.

3. 최적화 전략

• 손실 함수: 모델은 베이스라인이 포함된 정책 경사(policy gradient, REINFORCE) 방식을 사용하여 최적화됩니다. 손실은 토큰 단위 보상으로 분해되어, 마스크된 관측 하에서 특정 토큰을 생성하는 것의 이점(advantage)을 효율적으로 추정할 수 있게 합니다.
• 훈련 프로토콜:
  1. 사전 훈련 (Pre-training): 일부 토큰을 마스킹한 상태로 정답 데이터에 대해 사전 학습된 LM을 미세 조정(M-FT)하여, 모델이 마스크된 입력 형식에 적응하도록 합니다.
  2. RL 훈련: 결합 손실 $L = L_{MLE} + \lambda L_{RL}$을 최소화합니다. 여기서 $L_{MLE}$는 원래 데이터셋으로부터의 정책 드리프트(policy drift)를 방지하고, $L_{RL}$은 특정 목적 함수(예: ROUGE)를 최적화합니다.

주요 기여

1. 세미 오프라인 RL의 정의: 본 논문은 오프라인 및 온라인 RL 사이를 부드럽게 전환하는 이론적 설계 공간을 정의하여, 탐색 능력과 최적화 비용을 기준으로 서로 다른 설정들을 비교할 수 있는 토대를 제공합니다.
2. 이론적 최적성: 저자들은 최소 최적화 비용, 최소 점근적 편향, 최소 과적합 오차 경계를 만족하는 이론적으로 최적인 RL 설정(마스크된 관측)을 도출했습니다.
3. 효율적인 구현: 이 최적의 설정이 정적 데이터와 마스크 토큰을 혼합함으로써 구현될 수 있음을 입증하였으며, 이를 통해 단 1번의 순전파로 $K$개의 샘플을 탐색할 수 있음을 보여주었습니다.
4. 실험적 검증: 광범위한 실험을 통해 제안된 방법이 최신 기법들(온라인 Actor-Critic 및 오프라인 BRIO 포함)과 대등하거나 더 우수한 성능을 보이면서도 계산 비용을 크게 절감함을 입증했습니다.

실험 결과

방법론은 네 가지 텍스트 생성 작업(CNN/DM 및 SAMSum 요약, SQuAD 질문 생성, XSum 극단적 요약)에 대해 평가되었습니다.

• 성능 대 비용: 제안된 방법("OURS")은 최신 기법들(예: BRIO, AVG, SC, AC)과 대등하거나 더 나은 결과를 달성하면서도, $N \times T$ 또는 $N \times K \times T$번의 순전파가 필요한 멀티 샘플 온라인 방식과 달리 단 $N$번의 순전파(여기서 $N$은 인스턴스 수)만을 요구했습니다.
• 효율성: SQuAD 및 SAMSum 데이터셋에서 세미 오프라인 방식은 BRIO 및 AVG에 비해 가장 빠른 최적화 속도와 가장 낮은 훈련 시간 소비를 보여주었습니다.
• 절제 연구 (Ablation Studies):
  • 마스크 정보를 제거하면(Lemma 1 위반) 부정확한 보상 추정으로 인해 성능 저하가 발생했습니다.
  • 불필요한 정보를 추가하면(Lemma 2 위반) 과적합이 발생하여 테스트 세트 성능이 낮아졌습니다.
  • 탐색을 위해 "더 어려운" 정적 데이터셋(DATA-, 낮은 메트릭 후보군)을 사용하는 것이 고품질 후보군(DATA+)을 사용하는 것보다 더 나은 결과를 냈는데, 이는 하위 최적 지점의 주변부를 탐색하는 것이 개선 방향을 식별하는 데 도움이 된다는 것을 시사합니다.

의의 및 주장

본 논문은 세미 오프라인 RL이 이론적 토대를 제공하며, 온라인 RL의 높은 계산 비용에 대한 실질적인 해결책을 제시한다고 주장합니다. 마스크된 관측 전략을 활용함으로써, 이 방법은 대규모 사전 학습된 언어 모델이 오프라인 훈련의 효율성을 유지하면서도 복잡한 목적 함수(예: 인간 피드백 또는 ROUGE)를 위해 환경을 탐색하고 최적화할 수 있게 합니다. 저자들은 이 접근 방식이 순수 오프라인 방식보다 더 나은 개선 방향을 식별할 수 있게 하면서도, 전면적인 온라인 탐색의 과도한 시간 비용은 피할 수 있게 한다고 단언합니다.