Towards a theory of inference-time alignment with unknown rewards

기술적 요약: 미지의 보상을 가진 추론 시점 정렬(Inference-time Alignment) 이론을 향하여

문제 정의
본 연구는 생성 모델의 응답을 테스트 시점에 보상 신호에 따라 선택하거나 유도하는 패러다임인 **추론 시점 정렬(inference-time alignment)**의 통계적 학습 기초를 다룹니다. 이는 모델 파라미터를 업데이트하지 않는 방식(RLHF와는 다름)으로, 이 문제를 미지의 보상(unknown rewards) 가정하에서의 약한-강한 학습(weak-to-strong learning) 문제로 정식화합니다.

설정은 다음과 같습니다:

• 참조 정책 ($\pi_{ref}$): "약한 학습기"(예: 사전 학습된 또는 SFT 모델)로, 후보 응답을 생성합니다. 이 모델은 임의의 입력 $x$에 대해 "좋은" 응답(보상 $r^*(x, y) = 1$)을 생성할 확률 $\gamma$가 일정하다고 가정합니다.
• 보상 클래스 ($\mathcal{R}$): 이진 보상 함수 $r: \mathcal{X} \times \mathcal{Y} \to \{0, 1\}$의 집합입니다. 실제 보상 $r^*$은 알 수 없지만 $\mathcal{R}$ 내에 존재합니다.
• 데이터 생성: 훈련 데이터는 $m$개의 i.i.d. 샘플로 구성됩니다. 각 샘플은 입력 $x_i$와 $\pi_{ref}(\cdot|x_i)$로부터 추출된 $k$개의 응답 슬레이트(slate) $\{y_i^{(1)}, \dots, y_i^{(k)}\}$를 포함합니다. 학습자는 좋은 응답의 인덱스 집합 $A_{r^*}(z_i) = \{j : r^*(x_i, y_i^{(j)}) = 1\}$을 관찰합니다.
• 목표: 테스트 입력 $x$와 $\pi_{ref}$로부터 추출된 $k$개의 후보 슬레이트가 주어졌을 때, 높은 확률($1-\epsilon$)로 $r^*(x, y^{(j)}) = 1$을 만족하는 인덱스 $j$를 선택하는 학습 알고리즘을 설계하는 것입니다.

저자들은 사전에 추정된 보상 모델에 접근할 수 있다는 가정을 제거함으로써 기존의 이론적 문헌과 본 연구를 차별화합니다. 대신, 이들은 보상 함수 자체를 유한한 선호 데이터를 통해 처음부터 학습해야 하는 문제로 취급합니다.

방법론 및 핵심 개념
본 논문은 보상 클래스의 학습 가능성을 특징짓기 위해 새로운 조합적 차원을 도입합니다:

1. 정렬 차원 ($d_{ALN}(\mathcal{R})$):
   입력 집합 $\{x_1, \dots, x_d\}$가 $\mathcal{R}$에 의해 **정렬 파쇄(alignment shattered)**된다는 것은, 모든 이진 벡터 $b \in \{0, 1\}^d$에 대하여 $r_b(x_i, y_i^{(b_i)}) = 1$이고 $r_b(x_i, y_i^{(1-b_i)}) = 0$을 만족하는 보상 함수 $r_b \in \mathcal{R}$가 존재하는 쌍 $(y_i^{(0)}, y_i^{(1)})$이 존재함을 의미합니다. 정렬 차원은 이러한 파쇄된 집합의 최대 크기입니다.

2. 비지배 나타라잔 차원 (Non-Dominating Natarajan Dimension):
   학습 가능성을 증명하기 위해, 저자들은 부분 개념 클래스(partial concept classes)에 사용되는 One-Inclusion Graph (OIG) 알고리즘과 이 문제를 연결합니다. 이들은 유도된 집합-값 수용 클래스(set-valued acceptance class)에 대한 "비지배 나타라잔 차원"을 정의합니다. 또한, 이 차원이 정렬 차원에 의해 상한이 됨을 증명합니다.

3. 학습 알고리즘 (비지배 쌍에 대한 토너먼트):
   제안된 학습 절차는 전이적(transductive)이며 테스트 시점에 작동합니다:

   • 테스트 슬레이트 $Z$에 대해, 알고리즘은 어느 것도 다른 것의 부분 집합이 아닌($A \not\subseteq B$ 및 $B \not\subseteq A$) 비어 있지 않은 인덱스 부분 집합들의 모든 쌍 $(A, B)$를 고려합니다.
   • 각 쌍에 대해, 실제 수용 집합이 $A$이거나 $B$인 인스턴스를 포함하는 훈련 서브셋을 구성합니다.
   • 이 서브셋에 대해 OIG 알고리즘을 실행하여 실제 수용 집합이 $A$인지 $B$인지 예측합니다.
   • 이 과정은 생존한 집합들로 범위를 좁혀가는 "토너먼트"를 형성합니다. 최종 예측은 생존한 집합들의 교집합에 포함된 임의의 응답 인덱스입니다.

주요 결과
본 논문은 **정렬 학습 가능성(alignment learnability)**에 대한 완전한 특성화를 확립합니다:

• 정리 1: 보상 클래스 $\mathcal{R}$가 정렬 학습 가능하다는 것은 그 정렬 차원 $d_{ALN}(\mathcal{R})$이 유한하다는 것과 필요충분조건입니다.
• 상한 (정리 2): 만약 $d_{ALN}(\mathcal{R}) < \infty$라면, 제안된 토너먼트 알고리즘은 차원, $\epsilon$, $\delta$, $\gamma$의 유한한 함수인 샘플 수 $m$과 추론 시점 슬레이트 크기 $k$에 대해 $(\epsilon, \delta)$-학습 가능성을 달성합니다.
• 하한 (정리 3): 만약 $d_{ALN}(\mathcal{R}) = \infty$라면, 어떤 학습 알고리즘도 유한한 샘플로 정렬 학습 가능성을 달성할 수 없습니다. 증명은 확률적 방법(probabilistic method)을 사용하여 어려운 인스턴스를 구축하며, 무한한 차원이 학습자가 실제 보상 함수를 구별할 수 없는 분포의 존재를 함의함을 보여줍니다.

의의 및 주장
저자들은 이 연구를 일반적인 정렬 문제에 대한 엄격한 통계적 학습 이론을 위한 선구적인 단계로 위치시킵니다. 본 논문의 기여에 관한 주요 주장은 다음과 같습니다:

• 새로운 프레임워크: 보상이 미지이며 데이터로부터 완전히 학습되는 통계적 학습 문제로 정렬을 공식화한 최초의 연구로서, 양질의 보상 추정치에 접근할 수 있다고 가정하는 기존 연구들과 구별됩니다.
• 이론적 특성화: 정렬 차원의 도입은 이진 분류의 VC 차원과 유사하지만, 여러 후보 중 하나의 좋은 응답을 선택해야 하는 특정 제약에 맞춤화된, 학습 가능성에 대한 필요충분조건을 제공합니다.
• 기존 패러다임과의 차별화: 저자들은 본 설정이 멀티클래스 학습(하나의 레이블이 정답인 경우)이나 리스트 학습(레이블 리스트를 출력하는 경우)과 다르다는 점을 명확히 합니다. 본 설정에서는 여러 응답이 정답일 수 있지만, 학습자는 그중 하나를 성공적으로 식별해야 합니다.
• 범위의 겸손함: 저자들은 현재의 공식화가 이진 보상과 실현 가능한(realizable) 설정을 가정하고 있음을 인정하며 한계를 밝힙니다. 또한, 제안된 학습 절차가 표준적이지 않으며(테스트 시점 계산 및 튜플 순서 가정에 의존), 단일 응답 훈련 데이터, 실수 값 보상, 또는 선호 기반 학습(Bradley-Terry 모델)으로 이론을 확장하는 것이 향후 연구 과제로 남아 있음을 명시합니다.

본 논문은 정렬 차원을 근본적인 복잡도 척도로 확립하는 것이 언제 어떻게 추론 시점 정렬이 가능한지에 대한 이론적 토대를 제공하며, 잠재적으로 더 체계적인 원칙 개발을 가이드할 수 있다고 결론짓습니다.