라그랑주 승수와 KKT: 제약을 목적식 안으로 넣기

RLHF 학습 설정 파일에는 이런 줄이 있습니다.

kl_coef: 0.05          # 참조 모델에서 얼마나 멀어져도 되는가
target_kl: null

말로 하고 싶은 것은 「참조 모델에서 KL 기준으로 10 이내에 머물면서 보상을 최대로 받아라」입니다. 그런데 코드에 적힌 것은 예산 10이 아니라 계수 0.05입니다. 예산을 정하고 싶은데 손에 쥐어진 것은 가격이고, 그 가격을 얼마로 두면 예산이 얼마가 되는지는 돌려 봐야 압니다.

이 어긋남은 실수가 아니라 최적화의 구조 자체입니다. 제약이 붙은 문제를 제약 없는 문제로 바꾸는 방법이 딱 하나 있고, 바꾸는 순간 예산이 가격으로 변합니다. 그 방법이 라그랑주 승수이고, 이 글은 그 도구를 세웁니다. 세우고 나면 「KL 예산 안에서 보상을 최대화한다」를 식 한 줄로 적을 수 있고, 다음 글이 그 식을 실제로 풉니다.

제약 위의 정지 조건

접선 방향의 정지

먼저 등식 제약부터 봅니다. 함수 $f$ 를 최대화하되 $g(x)=0$ 을 만족하는 점들 중에서만 고른다고 합시다. 조건을 만족하는 점들의 모임을 실행 가능 영역이라 부르고, 여기서는 곡선 하나입니다.

최적점에서 무슨 일이 벌어지는지는 걸어 보면 압니다. 곡선 위의 어떤 점에 서서 곡선을 따라 조금 움직였을 때 $f$ 가 올라간다면, 거기는 최적점이 아닙니다 — 그쪽으로 가면 되니까요. 그러니 최적점에서는 곡선을 따라 어느 쪽으로 움직여도 $f$ 가 1차 근사로 변하지 않아야 합니다.

방향미분으로 적으면 곡선의 접선 방향 $v$ 에 대해 $\nabla f \cdot v = 0$ 입니다. 그런데 $g$ 는 곡선 위에서 값이 일정하므로 $\nabla g \cdot v = 0$ 도 성립합니다. 그래디언트가 등고선, 곧 함수값이 같은 점들을 이은 선에 수직이라는 그 성질이고, 제약 곡선 자체가 $g$ 의 값이 0인 등고선입니다. 평면에서 같은 방향에 수직인 벡터는 서로 평행하므로,

$\nabla f(x^\star) = \lambda\,\nabla g(x^\star)$

인 실수 $\lambda$ 가 있습니다. 이 $\lambda$ 를 라그랑주 승수라고 부릅니다. 두 그래디언트의 길이 비율이자 방향이 같은지 반대인지를 담는 수입니다.

그림으로 보면 더 분명합니다. $f$ 의 등고선이 제약 곡선을 가로지르면 곡선을 따라 더 높은 등고선으로 갈 수 있으니 아직 최적이 아니고, 접할 때 비로소 멈춥니다. 접한다는 것이 곧 두 그래디언트가 평행하다는 말입니다.

라그랑지안

이 조건을 함수 하나로 묶은 것이 라그랑지안입니다.

$\mathcal{L}(x, \lambda) = f(x) - \lambda\,g(x)$

$x$ 로 미분해 0으로 놓으면 $\nabla f = \lambda\nabla g$ 가 나오고, $\lambda$ 로 미분해 0으로 놓으면 $g(x)=0$ 이라는 제약이 그대로 나옵니다. 제약이 있는 문제가 변수를 하나 늘린 제약 없는 문제로 바뀐 것이고, 이것이 이 도구의 전부입니다. 원래 변수 $x$ 와 새로 들인 $\lambda$ 가 한 함수 안에 나란히 앉아, 앞쪽은 목적을 따르고 뒤쪽은 제약을 지키는 역할을 나눠 맡습니다.

최소화와 부호

교재를 몇 권 펼쳐 보면 라그랑지안이 $f - \lambda g$ 로도, $f + \lambda g$ 로도 적혀 있고 조건도 $\nabla f = \lambda\nabla g$ 와 $\nabla f = -\lambda\nabla g$ 로 갈립니다. 최소화 문제를 기본으로 두는 책은 흔히 $f + \lambda g$ 를 씁니다. 서로 다른 이론처럼 보이지만 같은 식입니다.

등식 제약에서 $\lambda$ 는 부호가 정해지지 않은 실수입니다. 한 책의 $\lambda$ 가 다른 책의 $-\lambda$ 일 뿐이라, 부호를 바꿔 적어도 풀리는 점은 똑같습니다. 최대화를 최소화로 바꾸는 것도 마찬가지로 $f$ 를 $-f$ 로 바꾸는 일이라 승수의 부호만 뒤집힙니다. 이 글은 최대화와 $f - \lambda g$ 를 기본으로 씁니다 — 맨 앞의 보상 최대화가 그 꼴이기 때문입니다. 부호가 뜻을 갖기 시작하는 것은 아래에서 부등식 제약을 만날 때입니다.

점으로 오그라든 제약

위의 논증은 「곡선을 따라 조금 움직인다」로 시작했습니다. 움직일 곡선이 없으면 어떻게 될까요. 평면에서 $g(x,y) = x^2 + y^2 = 0$ 을 제약으로 두면 이를 만족하는 점은 원점 하나뿐입니다. 실행 가능 영역이 점 하나로 오그라든 것입니다.

이 위에서 $f(x,y) = x$ 를 최대화하면 답은 당연히 원점입니다. 고를 점이 그것 하나니까요. 그런데 승수 조건을 세우면 $\nabla f = (1, 0)$ 이고 $\nabla g(0,0) = (0, 0)$ 이라 $(1,0) = \lambda\,(0,0)$ 을 만족하는 $\lambda$ 가 없습니다. 최적점은 있는데 승수가 없습니다.

어긋남의 원인은 첫 줄에 있습니다. 접선 방향이 하나도 없으니 「어느 방향으로 움직여도 $f$ 가 안 변한다」는 문장이 아무것도 요구하지 않고, 논증이 시작되지 않습니다. 증명에 숨어 있던 가정은 $\nabla g \ne 0$ 이었습니다 — 그래야 $g=0$ 이 제대로 된 곡선이 되고 접선이 생깁니다. 이 가정이 제약이 여럿일 때 어떤 모습이 되는지는 뒤에서 다시 봅니다.

단위원 위의 손계산

세 식과 두 해

단위원 위에서 $x+y$ 를 최대화해 보겠습니다.

$\max_{x,y}\ x + y \quad\text{s.t.}\quad x^2 + y^2 = 1$

라그랑지안을 적고 세 식을 세웁니다.

$\mathcal{L} = x + y - \lambda(x^2 + y^2 - 1)$

$$\begin{aligned} \frac{\partial\mathcal{L}}{\partial x} &= 1 - 2\lambda x = 0 &&\Longrightarrow\quad x = \frac{1}{2\lambda} \\ \frac{\partial\mathcal{L}}{\partial y} &= 1 - 2\lambda y = 0 &&\Longrightarrow\quad y = \frac{1}{2\lambda} \\ \frac{\partial\mathcal{L}}{\partial \lambda} &= -(x^2+y^2-1) = 0 &&\Longrightarrow\quad x^2+y^2 = 1 \end{aligned}$$

위 둘에서 $x=y$ 이고, 셋째에 넣으면 $2\cdot\frac{1}{4\lambda^2} = 1$ 이라 $\lambda^2 = \frac12$ 입니다. 해가 둘입니다. $\lambda = \frac{1}{\sqrt2}$ 를 되돌리면

$x^\star = y^\star = \frac{1}{\sqrt2} \approx 0.7071,\qquad f^\star = \sqrt{2} \approx 1.4142$

여기서 $\lambda = 0.7071$ 이라는 값이 그냥 계산 과정에서 나온 부산물처럼 보이는데, 아래 「그림자 가격」 절에서 보듯 이 숫자에도 뜻이 있습니다.

둘째 해와 최소점

남은 해 $\lambda = -\frac{1}{\sqrt2}$ 를 되돌리면 $x = y = -\frac{1}{\sqrt2}$ 이고 $f = -\sqrt2 \approx -1.4142$ 입니다. 이 점도 세 식을 모두 만족합니다. 그런데 값이 가장 작습니다 — 원 위에서 $x+y$ 가 가장 작은 점, 곧 최소점입니다.

라그랑주 조건은 「멈춘 점」을 찾는 조건이지 「최대인 점」을 찾는 조건이 아닙니다. 곡선을 따라 움직여도 $f$ 가 변하지 않는 자리라면 꼭대기든 바닥이든 다 걸립니다. 제약 없는 문제에서 도함수가 0인 점이 극대일 수도 극소일 수도 있는 것과 같은 사정입니다.

그래서 마지막 걸음은 후보마다 $f$ 값을 넣어 보는 일입니다. 원은 닫혀 있고 끝이 없는 곡선이라 그 위의 연속함수는 최대와 최소를 반드시 가지며, 둘 다 멈춘 점 가운데 있어야 합니다. 후보가 둘이니 큰 쪽 $\sqrt2$ 가 최대, 작은 쪽 $-\sqrt2$ 가 최소입니다.

두 접점의 승수 부호

두 해의 승수 부호가 다른 데도 뜻이 있습니다. $\nabla f = (1,1)$ 은 어디서나 같고, $\nabla g = (2x, 2y)$ 는 원 밖을 향합니다. 최대점 $\left(\frac{1}{\sqrt2}, \frac{1}{\sqrt2}\right)$ 에서는 $\nabla g = (\sqrt2, \sqrt2)$ 라 $\nabla f$ 와 같은 방향이고 $\lambda > 0$ 입니다. 최소점에서는 $\nabla g = (-\sqrt2, -\sqrt2)$ 라 반대 방향이고 $\lambda < 0$ 입니다.

읽으면 이렇습니다. 최대점에서 $f$ 를 더 키우는 방향은 원 밖이고, 최소점에서 $f$ 를 더 키우는 방향은 원 안입니다. 제약이 등식일 때는 안쪽이든 바깥쪽이든 못 가기는 마찬가지라 부호가 어느 쪽이어도 괜찮았습니다. 제약이 「원 안이면 된다」로 바뀌면 사정이 달라집니다. 최소점에서는 안으로 들어가면 되니 더 이상 멈춘 자리가 아니고, 최대점만 여전히 막혀 있습니다. 부등식 제약에서 승수의 부호가 조건이 되는 이유가 이 그림 안에 이미 들어 있습니다.

부등식 제약과 KKT

느슨한 제약과 팽팽한 제약

실제 문제는 대개 「같아야 한다」가 아니라 「이하여야 한다」입니다. KL 예산도 「KL이 정확히 10」이 아니라 「KL이 10 이하」입니다.

$\max_x\ f(x) \quad\text{s.t.}\quad g(x) \le c$

이때는 경우가 둘로 갈립니다.

• 제약이 느슨한 경우. 제약 없이 푼 최적점이 이미 $g(x) \le c$ 를 만족하면, 제약은 있으나 마나입니다. 이때 승수는 0입니다.
• 제약이 팽팽한 경우. 제약 없는 최적점이 밖에 있으면 최적점은 경계 $g(x)=c$ 위로 밀려납니다. 이때는 등식 제약과 똑같은 상황이고 승수가 0보다 큽니다.

간단한 예로 확인해 보겠습니다. $f(x) = -(x-3)^2$ 를 $x \le c$ 아래에서 최대화합니다. $c \ge 3$ 이면 제약 없는 최적점 $x=3$ 이 실행 가능 영역 안에 있으므로 $x^\star = 3$, $f^\star = 0$, $\lambda = 0$ 입니다. $c < 3$ 이면 $f$ 가 $x=3$ 까지 계속 오르므로 갈 수 있는 데까지 갑니다. $x^\star = c$, $f^\star = -(c-3)^2$ 이고 승수는 $\lambda = f'(c) = 2(3-c) > 0$ 입니다.

KKT 네 조건

두 경우를 한꺼번에 적은 것이 KKT 조건입니다 — Karush, Kuhn, Tucker 세 사람의 이름을 딴 조건 묶음이고, 라그랑주 승수를 부등식 제약으로 넓힌 것입니다. $\mathcal{L}(x,\lambda) = f(x) - \lambda(g(x)-c)$ 에 대해 최적점은 네 조건을 모두 만족합니다.

조건	식	뜻
정상성	$\nabla f(x^\star) = \lambda\,\nabla g(x^\star)$	라그랑지안의 그래디언트가 0
원시 실행 가능성	$g(x^\star) \le c$	답이 제약을 지킨다
쌍대 실행 가능성	$\lambda \ge 0$	가격이 음수일 수 없다
상보 여유	$\lambda\,(g(x^\star)-c) = 0$	둘 중 하나는 반드시 0

마지막 줄이 두 경우를 하나로 묶는 장치입니다. 상보 여유는 「제약이 남아돌면 그 제약의 가격은 0이고, 가격이 0이 아니면 제약은 남아돌지 않는다」는 조건입니다. 위 예에서 $c > 3$ 이면 $\lambda = 0$ 이라 왼쪽 인수가 0이고, $c < 3$ 이면 $x^\star = c$ 라 오른쪽 인수가 0입니다. 곱이 0이라는 한 줄이 경우 나눔을 그대로 담고 있습니다.

셋째 줄에도 이유가 있습니다. $g \le c$ 라는 제약에서 $\lambda$ 가 음수이면 라그랑지안이 $g$ 를 키우는 쪽에 상을 주게 되어, 제약을 어기는 방향으로 답을 밀어냅니다. 벌점이어야 할 항이 보상이 되어 버리는 것입니다. 앞 절의 원으로 말하면 음의 승수를 가진 최소점을 「원 안이면 된다」 문제의 답으로 받아들이는 셈입니다.

두 인수가 함께 0인 경계

상보 여유는 둘 중 하나가 0이라고만 말하고, 둘 다 0인 것을 막지 않습니다. 위 예에서 $c = 3$ 이 그 자리입니다. 제약 없는 최적점이 정확히 경계에 놓여 $x^\star = 3 = c$ 로 제약이 팽팽하면서도, 경계가 없었어도 거기서 멈췄을 것이라 $\lambda = 2(3-3) = 0$ 입니다. 제약이 닿아 있기는 한데 아무 힘도 쓰지 않는 상태입니다.

이 자리는 불안정합니다. $c$ 를 조금만 올리면 제약이 느슨해지고, 조금만 내리면 팽팽해집니다. 어느 제약이 일하고 있는지가 $c$ 의 아주 작은 흔들림에 따라 바뀌는 것입니다. 식으로 보면 $\lambda(c) = \max(0,\ 2(3-c))$ 와 $x^\star(c) = \min(c,\ 3)$ 이 둘 다 $c = 3$ 에서 꺾입니다. 값은 이어지지만 $\lambda$ 의 기울기는 $-2$ 에서 0으로, $x^\star$ 의 기울기는 1에서 0으로 바뀌어 그 점에서 미분이 없습니다. 「지금 어떤 제약이 팽팽한가」를 추측하며 나아가는 풀이법이 이런 자리 근처에서 추측을 번갈아 뒤집으며 헤매는 이유가 이것입니다.

부호를 정하는 네 조합

최대화냐 최소화냐, 제약이 $\le$ 냐 $\ge$ 냐로 네 조합이 나오고, 교재마다 라그랑지안의 부호가 달라 보이는 또 하나의 이유가 여기에 있습니다. 규칙은 하나입니다 — 제약을 어기는 쪽으로 가면 목적이 나빠지도록 부호를 고르고, 그렇게 고르면 승수는 언제나 0 이상입니다.

문제	제약	라그랑지안	정상성
최대화	$g \le c$	$f - \lambda(g - c)$	$\nabla f = \lambda\nabla g$
최대화	$g \ge c$	$f + \lambda(g - c)$	$\nabla f = -\lambda\nabla g$
최소화	$g \le c$	$f + \lambda(g - c)$	$\nabla f = -\lambda\nabla g$
최소화	$g \ge c$	$f - \lambda(g - c)$	$\nabla f = \lambda\nabla g$

정상성 열을 말로 읽으면 네 줄이 같은 문장입니다. 경계에서 멈춘 점에서는 $f$ 를 좋게 하는 방향이 정확히 제약을 어기는 방향을 가리킵니다. 최대화에서 좋게 하는 방향은 $\nabla f$, 최소화에서는 $-\nabla f$ 이고, $g \le c$ 를 어기는 방향은 $\nabla g$, $g \ge c$ 를 어기는 방향은 $-\nabla g$ 입니다. 그 둘이 같은 쪽을 가리킨다는 것이 $\lambda \ge 0$ 입니다.

제약이 여럿인 문제

등식 제약 여럿

제약이 $g_1(x) = 0, \dots, g_m(x) = 0$ 으로 $m$ 개면 승수도 $m$ 개이고 조건은

$\nabla f(x^\star) = \sum_{i=1}^{m} \lambda_i\,\nabla g_i(x^\star)$

입니다. 논리는 하나일 때와 같습니다. 실행 가능 영역은 모든 제약을 동시에 지키는 점들이고, 그 위의 접선 방향 $v$ 는 모든 $\nabla g_i$ 에 수직입니다. 최적점에서 $\nabla f$ 도 그런 $v$ 전부에 수직이어야 하므로, $\nabla f$ 는 $\nabla g_i$ 들이 펼치는 공간 안에 있어야 합니다. 그 공간 안의 벡터를 $\nabla g_i$ 들의 조합으로 적은 계수가 승수들입니다.

예로 3차원에서 $x+y+z$ 를 최대화하되 $x^2+y^2+z^2 = 1$ 과 $z = 0$ 을 함께 지키게 합시다. 두 제약을 같이 지키는 점들은 바닥 평면 위의 단위원이라 답은 앞 절과 같은 $\left(\frac{1}{\sqrt2}, \frac{1}{\sqrt2}, 0\right)$ 입니다. 거기서 $\nabla f = (1,1,1)$ 이고 두 제약의 그래디언트는 $(\sqrt2, \sqrt2, 0)$ 과 $(0, 0, 1)$ 이므로

$(1, 1, 1) = \frac{1}{\sqrt2}\,(\sqrt2, \sqrt2, 0) + 1\cdot(0, 0, 1)$

로 $\lambda_1 = \frac{1}{\sqrt2}$, $\lambda_2 = 1$ 입니다. 첫째 승수는 앞의 원 문제에서 구한 값 그대로이고, 둘째 승수는 $z$ 방향으로 올라가고 싶은 힘을 $z = 0$ 이 막고 있는 크기입니다.

활성 집합

부등식이 섞이면 제약마다 느슨하거나 팽팽합니다. 최적점에서 등호로 성립하는 제약, 곧 팽팽한 제약을 활성 제약이라 부르고 그 모임을 활성 집합이라 합니다. KKT 조건은 제약마다 상보 여유를 따로 걸기 때문에, 활성이 아닌 제약은 승수가 0이 되어 정상성 식에서 저절로 빠집니다. 결국 최적점에서는 활성 제약만 등식 제약처럼 살아 있습니다.

원 문제에 제약을 하나 더 얹어 보겠습니다.

$\max_{x,y}\ x + y \quad\text{s.t.}\quad x^2 + y^2 \le 1,\quad x \le \tfrac12$

둘째 제약이 없을 때의 답은 $x = 0.7071$ 이라 $x \le \frac12$ 를 어깁니다. 그러니 둘째 제약은 활성입니다. 원을 따라 $x$ 를 줄이면 $x + \sqrt{1-x^2}$ 는 $x < \frac{1}{\sqrt2}$ 구간에서 $x$ 와 함께 커지므로, 허락된 가장 큰 $x = \frac12$ 로 가서 $y = \frac{\sqrt3}{2}$, $f^\star = \frac12 + \frac{\sqrt3}{2} \approx 1.3660$ 이 됩니다. 두 제약이 모두 활성입니다.

정상성은 $(1, 1) = \lambda_1\,(1, \sqrt3) + \lambda_2\,(1, 0)$ 이고, 둘째 성분에서 $\lambda_1 = \frac{1}{\sqrt3} \approx 0.5774$, 첫째 성분에서 $\lambda_2 = 1 - \frac{1}{\sqrt3} \approx 0.4226$ 입니다. 둘 다 0 이상이니 KKT를 만족합니다. 둘째 제약을 $x \le 0.9$ 로 바꾸면 그것은 비활성이 되어 $\lambda_2 = 0$ 이고, 답은 원 문제의 답으로 돌아갑니다.

기하로 보면 부호 조건이 모양을 갖습니다. 활성 제약의 그래디언트들에 0 이상의 계수를 붙여 만들 수 있는 벡터 전체는 꼭짓점에서 뻗어 나가는 부채꼴, 곧 원뿔을 이룹니다. KKT의 정상성과 $\lambda_i \ge 0$ 을 합치면 「$\nabla f$ 가 그 원뿔 안에 있다」는 한 문장이 됩니다. 원뿔 밖을 가리키면 두 제약을 다 지키면서 $f$ 를 키울 틈이 남아 있다는 뜻입니다.

제약 자격조건

「제약이 점으로 오그라든 예」에서 승수가 없었던 이유를 이제 일반형으로 적을 수 있습니다. 여러 벡터 중 하나를 나머지의 조합으로 적을 수 있으면 그 벡터들이 일차종속이라 하고, 그럴 수 없으면 일차독립이라 합니다. 최적점에서 활성 제약의 그래디언트들이 일차독립이면 승수가 반드시 있고 하나로 정해집니다. 라그랑주 조건이나 KKT 조건을 믿고 써도 되게 해 주는 이런 가정을 제약 자격조건이라 부르고, 이것이 그중 가장 흔한 형태입니다. 제약이 하나일 때는 「$\nabla g \ne 0$」 이 곧 이 조건입니다.

조건이 깨지는 모습은 두 가지입니다. 하나는 승수가 여럿이 되는 경우입니다. 위 문제에 $x \le \frac12$ 와 뜻이 같은 $2x \le 1$ 을 한 번 더 적으면 두 그래디언트 $(1,0)$ 과 $(2,0)$ 이 일차종속이 되고, 정상성은 $\lambda_2 + 2\lambda_3 = 1 - \frac{1}{\sqrt3}$ 만 요구합니다. 이 합을 만족하는 짝은 무수히 많습니다. 답은 하나인데 가격이 두 제약 사이에 어떻게 나뉘는지가 정해지지 않은 것입니다.

다른 하나는 승수가 아예 없는 경우입니다. 원판 $x^2 + y^2 \le 1$ 과 $x \ge 1$ 을 함께 걸면 둘을 동시에 지키는 점은 $(1, 0)$ 하나뿐이고, 두 경계는 거기서 가로지르지 않고 접합니다. $x \ge 1$ 을 $1 - x \le 0$ 으로 고쳐 적으면 두 그래디언트 $(2, 0)$ 과 $(-1, 0)$ 이 한 직선 위에 있어 원뿔이 가로축으로 납작해지므로, 목적이 $f = y$ 라면 $\nabla f = (0, 1)$ 을 그 조합으로 만들 방법이 없습니다. 최적점은 있는데 KKT 조건을 만족하는 승수가 없습니다. 첫 절의 점 제약과 같은 병입니다.

그림자 가격

최적값의 기울기

$\lambda$ 가 부산물이 아니라는 것을 이제 보겠습니다. 예산 $c$ 를 조금 늘려 주면 최적값 $f^\star$ 가 얼마나 오를까요.

$\frac{d f^\star(c)}{dc} = \lambda$

승수는 제약을 한 단위 풀었을 때 얻는 이득입니다. 경제학에서 이 값을 그림자 가격이라 부릅니다 — 예산 한 단위를 사 올 수 있다면 그 값어치가 얼마인가를 뜻합니다.

앞의 예에서 확인됩니다. $f(x) = -(x-3)^2$ 문제에서 $c<3$ 일 때 $f^\star(c) = -(c-3)^2$ 이므로 미분하면 $-2(c-3) = 2(3-c)$ 이고, 이것이 앞에서 구한 $\lambda$ 와 같습니다. 원 문제도 마찬가지입니다. 제약을 $x^2+y^2 = c$ 로 두면 최적값이 $f^\star(c) = \sqrt{2c}$ 이고,

$\frac{d f^\star}{dc} = \frac{1}{\sqrt{2c}} \;\Big|_{c=1} = \frac{1}{\sqrt2} \approx 0.7071$

앞에서 구한 $\lambda$ 와 소수점까지 같습니다.

예산과 가격의 쌍대

이 해석이 맨 앞의 어긋남을 설명합니다. 예산 $c$ 를 정하는 것과 가격 $\lambda$ 를 정하는 것은 같은 문제의 앞뒤 면입니다. 예산을 정하면 그에 맞는 가격이 따라 나오고, 가격을 정하면 그에 맞는 예산이 따라 나옵니다. 이렇게 짝을 이루는 두 문제를 서로의 쌍대 문제라고 부릅니다.

다만 방향이 다릅니다. 예산을 정해 놓고 가격을 찾으려면 안쪽 문제를 풀면서 바깥에서 $\lambda$ 를 조절해야 하는데, 이것은 학습 루프 안에 루프를 하나 더 두는 일입니다. 반대로 가격을 고정하면 벌점이 붙은 목적식 하나를 그냥 최대화하면 됩니다. 그래서 설정 파일에 예산이 아니라 계수가 적혀 있는 것입니다.

꺾이는 최적값

「예산에 맞는 가격이 따라 나온다」는 말에는 단서가 붙습니다. 최적값 $f^\star(c)$ 가 매끄러울 때만 그 기울기가 하나로 정해집니다. 활성 집합이 바뀌는 $c$ 에서는 $f^\star$ 가 꺾일 수 있습니다.

가장 단순한 예로 $x$ 를 최대화하되 예산 $x \le c$ 와 별도의 한계 $x \le 2$ 를 함께 겁니다. 최적값은 $f^\star(c) = \min(c,\ 2)$ 입니다. $c < 2$ 이면 예산이 활성이고 예산의 가격은 1이며, $c > 2$ 이면 한계가 대신 막고 있어 예산은 비활성이라 가격이 0입니다. $c = 2$ 에서는 왼쪽 기울기가 1, 오른쪽 기울기가 0으로 갈려 미분이 없습니다.

그 자리에서 KKT를 세우면 두 제약이 모두 활성이고 그래디언트가 둘 다 1이라, 앞 절의 「같은 제약을 두 번 적은」 경우가 됩니다. 정상성은 $\lambda_1 + \lambda_2 = 1$ 만 요구하고, 예산의 가격 $\lambda_1$ 은 0과 1 사이 어느 값이든 됩니다. 꺾인 점에서 기울기가 하나로 안 정해진다는 것과 승수가 하나로 안 정해진다는 것은 같은 사실의 두 얼굴입니다.

예산 방식의 흔들림

이 꺾임이 두 방식의 차이를 드러냅니다. 가격을 고정하는 쪽에서 보면, 0과 1 사이의 가격은 어느 것이든 같은 자리 $x = 2$ 로 데려갑니다. 가격을 조금 틀리게 골라도 답이 움직이지 않습니다.

예산을 고정하는 쪽은 사정이 반대입니다. 예산을 2 바로 아래에 두면 필요한 가격이 1이고, 2 바로 위에 두면 0입니다. 예산이 조금만 움직여도 맞춰야 할 가격이 한 번에 1만큼 뜁니다. 바깥 루프가 KL 같은 측정값을 보며 가격을 되먹임으로 찾아가는 방식이라면, 측정값이 이 경계 근처에서 흔들릴 때마다 가격이 한쪽 끝에서 다른 쪽 끝으로 끌려다닙니다. 학습 중에 계수가 튀는 모습으로 나타나는 것이 이것입니다.

가격을 직접 정하면 이 꺾임을 밟을 일이 없습니다. 되먹임이 없으니 가격은 처음 정한 값에 머물고, 예산은 그 가격이 데려가는 곳에서 결과로 정해집니다. 두 방식이 쌍대로 같은 문제를 가리키더라도, 어느 쪽을 손잡이로 쥐느냐에 따라 학습의 안정성이 달라집니다.

KL 예산 문제의 라그랑지안

벌점 항의 정체

이제 맨 앞의 문장을 식으로 옮길 수 있습니다. 하고 싶은 말은 이것입니다.

$\max_{\pi}\ \mathbb{E}_{y\sim\pi}\big[r(y)\big] \quad\text{s.t.}\quad \mathrm{KL}\big(\pi \,\|\, \pi_{\text{ref}}\big) \le \varepsilon$

부등식 제약이므로 KKT를 씁니다. 라그랑지안은

$\mathcal{L}(\pi, \beta) = \mathbb{E}_{y\sim\pi}\big[r(y)\big] - \beta\Big(\mathrm{KL}\big(\pi\,\|\,\pi_{\text{ref}}\big) - \varepsilon\Big)$

이고, $\beta\varepsilon$ 은 $\pi$ 와 무관한 상수라 최대화에 영향을 주지 않으므로 떼어내면

$\max_{\pi}\ \mathbb{E}_{y\sim\pi}\big[r(y)\big] - \beta\,\mathrm{KL}\big(\pi\,\|\,\pi_{\text{ref}}\big)$

가 남습니다. 앞선 글에서 벌점 항으로 그냥 붙여 두었던 그 식이고, 설정 파일의 kl_coef가 바로 $\beta$ 입니다. 벌점처럼 보이던 항이 사실은 예산 제약의 라그랑주 승수였던 것입니다.

읽는 법도 따라옵니다. $\beta$ 는 「KL 1 단위의 가격」입니다. 크게 두면 KL이 비싸지므로 정책이 참조 모델 근처에 머물고, 작게 두면 싸지므로 멀리 나갑니다. 그리고 상보 여유가 말해 주는 것도 있습니다 — $\beta$ 를 충분히 작게 두어 KL이 예산보다 한참 아래에 머문다면 그 제약은 아무 일도 하고 있지 않은 것이고, 그 자리에서는 $\beta$ 를 더 줄여도 학습이 거의 달라지지 않습니다.

정규화 승수

위 식에는 적지 않은 제약이 하나 숨어 있습니다. $\pi$ 는 응답 $y$ 마다 확률을 매기는 분포이므로 모든 응답의 확률을 더하면 1이어야 합니다. 이것도 등식 제약이고, 승수가 하나 더 붙습니다. 그 승수를 $\mu$ 라 하면 라그랑지안은

$\mathcal{L}(\pi, \beta, \mu) = \sum_y \pi(y)\,r(y) - \beta\Big(\sum_y \pi(y)\log\frac{\pi(y)}{\pi_{\text{ref}}(y)} - \varepsilon\Big) - \mu\Big(\sum_y \pi(y) - 1\Big)$

가 됩니다. 승수가 둘입니다 — 예산의 가격 $\beta$ 와 정규화의 가격 $\mu$ 입니다. 앞 절의 말로 하면 등식 제약 하나와 부등식 제약 하나가 함께 걸린 문제입니다.

응답 하나의 확률 $\pi(y)$ 로 미분해 0으로 놓으면

$r(y) - \beta\Big(\log\frac{\pi(y)}{\pi_{\text{ref}}(y)} + 1\Big) - \mu = 0$

이고, 이 식을 $\pi(y)$ 에 대해 풀면 $\mu$ 는 모든 $y$ 에 공통으로 곱해지는 상수 자리에 들어갑니다. 그 상수를 확률의 합이 1이 되도록 맞추는 일이 곧 $\mu$ 를 정하는 일이고, 다음 글의 닫힌 해가 바로 이 정규화에서 나옵니다. 확률이 0 이상이어야 한다는 제약도 있지만, 로그가 들어 있어 해가 0에 닿지 않으므로 그 제약들은 전부 비활성이고 승수가 0이라 식에서 빠집니다.

적응형 KL 계수

맨 앞 설정에는 target_kl 이라는 칸도 있었습니다. 예산을 직접 정하고 싶을 때 쓰는 자리이고, 이 경우 $\beta$ 는 고정된 값이 아니라 학습 중에 조절되는 변수가 됩니다. 조절 규칙은 라그랑지안을 $\beta$ 에 대해 미분한 데서 나옵니다.

$\frac{\partial \mathcal{L}}{\partial \beta} = -\big(\mathrm{KL} - \varepsilon\big)$

우리는 $\beta$ 에 대해서는 라그랑지안을 최소화해야 하므로 이 미분의 반대 방향으로 갑니다.

$\beta \leftarrow \beta + \eta\,\big(\mathrm{KL} - \varepsilon\big)$

읽으면 자명합니다 — KL이 예산을 넘고 있으면 가격을 올리고, 예산 아래에 머물면 가격을 내립니다. 이것을 쌍대 상승법이라 부르고, RLHF 구현에서 흔히 「적응형 KL 계수」라고 부르는 것이 정확히 이 갱신입니다. 그리고 $\beta$ 가 0 아래로 내려가지 않게 막는 코드가 붙어 있다면 그것은 KKT의 셋째 조건 $\lambda \ge 0$ 을 지키는 장치입니다. 앞 절에서 본 흔들림이 생기는 자리가 바로 이 갱신입니다.

쌍대 함수의 볼록성

그런데 왜 $\beta$ 에 대해서는 최소화일까요. 가격을 하나 정하고 그 가격에서 $\pi$ 로 라그랑지안을 최대화한 값을 $d(\beta) = \max_\pi \mathcal{L}(\pi, \beta)$ 라 두고, 이것을 쌍대 함수라 부릅니다. 어떤 가격에서든 $d(\beta)$ 는 원래 문제의 최적값보다 작지 않습니다 — 예산을 지키는 정책에서는 $\beta(\mathrm{KL} - \varepsilon)$ 이 0 이하라 라그랑지안이 원래 목적 이상이기 때문입니다. 그러니 가장 빡빡한 위쪽 경계를 주는 가격, 곧 $d$ 를 가장 작게 만드는 $\beta$ 를 찾는 것이 옳은 가격을 찾는 일입니다.

이 함수는 $\beta$ 에 대해 볼록합니다. 볼록 함수란 그래프 위 두 점을 이은 선분이 그래프 아래로 내려가지 않는 함수, 곧 기울기가 줄어들지 않는 함수입니다. 근거는 짧습니다. $\pi$ 를 하나 고정하면 $\mathcal{L}(\pi, \beta)$ 는 $\beta$ 에 대한 일차식, 곧 직선이고, $d$ 는 그런 직선들 가운데 가장 높은 것을 매 $\beta$ 마다 고른 윗면입니다. 직선들의 윗면은 언제나 볼록합니다.

볼록이면 기울기가 오른쪽으로 갈수록 커집니다. 그리고 최적 정책 $\pi_\beta$ 에서는 $\pi$ 로의 미분이 0이므로 $d$ 를 $\beta$ 로 미분하면 $\mathcal{L}$ 을 $\beta$ 로 직접 미분한 항만 남아

$d'(\beta) = \varepsilon - \mathrm{KL}\big(\pi_\beta \,\|\, \pi_{\text{ref}}\big)$

입니다. 이 기울기가 $\beta$ 와 함께 커진다는 것은 $\mathrm{KL}(\pi_\beta)$ 가 $\beta$ 와 함께 줄어든다는 뜻입니다. 가격을 올리면 KL이 단조로 줄어든다는 직관이 볼록성 한 줄에서 나옵니다. 그리고 $d$ 의 최소점에서는 기울기가 0이라 KL이 정확히 $\varepsilon$ 이 됩니다 — 적응형 계수가 찾아가는 자리가 그곳입니다.

응답이 셋뿐인 작은 예로 확인합니다. 최적 정책의 꼴은 다음 글에서 유도할 결과를 미리 빌려 씁니다.

import numpy as np

p_ref = np.array([0.5, 0.3, 0.2])   # 참조 정책
r = np.array([1.0, 2.0, 0.0])       # 응답별 보상
eps = 0.2                           # KL 예산

def solve(beta):
    w = p_ref * np.exp(r / beta)
    pi = w / w.sum()                 # 합이 1이 되게 정규화
    kl = np.sum(pi * np.log(pi / p_ref))
    d = pi @ r - beta * (kl - eps)   # 쌍대 함수 d(β)
    return kl, d

for beta in [0.25, 0.5, 1.0, 2.0, 4.0]:
    kl, d = solve(beta)
    print(f"beta={beta:<5} KL={kl:.3f}  d={d:.3f}")

beta=0.25  KL=1.053  d=1.757
beta=0.5   KL=0.587  d=1.605
beta=1.0   KL=0.205  d=1.529
beta=2.0   KL=0.058  d=1.619
beta=4.0   KL=0.015  d=1.961

$\beta$ 를 올릴수록 KL이 1.053에서 0.015까지 한 번도 거꾸로 가지 않고 줄어듭니다. 쌍대 함수는 내려갔다가 다시 올라가는 그릇 모양이고, 바닥 근처인 $\beta = 1$ 에서 KL이 0.205로 예산 0.2에 가장 가깝습니다. 예산 0.2를 지키고 싶다면 가격은 1 근처라는 것을 쌍대 함수의 바닥이 알려 줍니다.

여기까지가 문제를 적는 일입니다. 실제로 저 $\max$ 를 푸는 일, 즉 최적 정책 $\pi^\star$ 를 닫힌 형태로 구하는 일이 다음 글의 주제입니다. 그리고 그 답을 뒤집어 보상을 정책으로 표현하면, 지난 글의 Bradley-Terry 손실에 그대로 대입할 수 있게 됩니다. 맨 앞의 kl_coef: 0.05 한 줄은 이제 이렇게 읽힙니다 — KL 한 단위를 보상 0.05만큼의 값으로 치겠다는 가격 선언이고, 그 가격이 데려가는 예산은 학습이 끝나 봐야 압니다.

정리

• 등식 제약 $g(x)=0$ 아래에서 $f$ 를 최적화하면, 최적점에서 곡선을 따라 움직여도 $f$ 가 변하지 않아야 하므로 $\nabla f = \lambda\nabla g$ 가 성립한다. 이 $\lambda$ 가 라그랑주 승수이고, 기하적으로는 등고선이 제약 곡선에 접한다는 뜻이다. 최소화로 적거나 $f+\lambda g$ 로 적으면 부호만 뒤집힌다.
• 라그랑지안 $\mathcal{L} = f - \lambda g$ 를 두면 $x$ 로 미분한 것이 위 조건, $\lambda$ 로 미분한 것이 제약이다. 제약이 있는 문제가 변수 하나 늘린 제약 없는 문제로 바뀐다.
• 단위원 위에서 $x+y$ 의 멈춘 점은 둘이다. $\lambda=1/\sqrt2$ 인 점이 최대 $\sqrt2$, $\lambda=-1/\sqrt2$ 인 점이 최소 $-\sqrt2$ 이고, 어느 쪽인지는 값을 넣어 가른다.
• 부등식 제약에서는 KKT 조건 넷이 최적점을 특징짓는다 — 정상성, 원시 실행 가능성 $g\le c$, 쌍대 실행 가능성 $\lambda\ge0$, 상보 여유 $\lambda(g-c)=0$. 두 인수가 함께 0인 경계는 활성 여부가 작은 흔들림에 뒤집히는 불안정한 자리다.
• 제약이 여럿이면 $\nabla f = \sum\lambda_i\nabla g_i$ 이고, 활성 제약만 살아남는다. 부등식에서는 $\nabla f$ 가 활성 그래디언트들의 원뿔 안에 있어야 한다. 활성 그래디언트가 일차종속이면 제약 자격조건이 깨져 승수가 여럿이거나 아예 없다.
• 승수는 그림자 가격이다 — $df^\star/dc = \lambda$ 로, 예산을 한 단위 풀었을 때 최적값이 오르는 폭이다. 활성 집합이 바뀌는 자리에서는 $f^\star$ 가 꺾여 가격이 하나로 안 정해지고, 예산을 되먹임으로 맞추는 방식은 거기서 가격이 튄다.
• 「KL 예산 안에서 보상을 최대화한다」의 라그랑지안은 $\mathbb{E}_\pi[r] - \beta(\mathrm{KL} - \varepsilon)$ 이고, 상수를 떼면 흔히 보는 $\mathbb{E}_\pi[r] - \beta\,\mathrm{KL}(\pi\|\pi_{\text{ref}})$ 다. kl_coef가 곧 승수 $\beta$ 이고, 확률의 합이 1이라는 제약에 둘째 승수 $\mu$ 가 붙는다.
• 예산을 직접 정하고 싶으면 $\beta \leftarrow \beta + \eta(\mathrm{KL}-\varepsilon)$ 로 가격을 갱신한다(쌍대 상승법). 쌍대 함수가 $\beta$ 에 대해 볼록이라 $\beta$ 를 올리면 KL이 단조로 줄고, 쌍대 함수의 바닥에서 KL이 예산과 같아진다.

---

읽어주셔서 감사합니다. 😊