💬 NLP

Safe Inference-Time Alignment via Lagrangian Reward Augmentation

이 논문은 재학습 없이 안전 제약 조건을 강제하기 위해 이중 변수를 통해 증강된 보상 신호를 동적으로 보정함으로써, 유용성-무해성 트레이드오프를 개선하고 파인튜닝 기반 방법의 성능에 근접하는 일반적인 추론 시 정렬 프레임워크인 라그랑주 보상 증강(LARA)을 제안한다.

원저자: Yaswanth Chittepu, Ativ Joshi, Sohini Chintala, Scott Niekum

게시일 2026-07-07

📖 4 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Yaswanth Chittepu, Ativ Joshi, Sohini Chintala, Scott Niekum

원본 논문은 CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/) 라이선스로 제공됩니다. ✨ 이것은 아래 논문에 대한 AI 생성 설명입니다. 저자가 작성하거나 승인한 것이 아닙니다. 기술적 정확성을 위해서는 원본 논문을 참조하세요. 전체 면책 조항 읽기

LARA(Lagrangian Reward Augmentation) 논문 설명: 쉽고 창의적인 비유를 곁들인 해설

거대한 문제: "유능하지만 위험한" 로봇

당신에게 매우 똑똑한 로봇 비서(거대 언어 모델)가 있다고 상상해 보세요. 당신은 이 로봇이 **유능(helpful)**하기를 바라지만(좋은 답변을 제공함), 동시에 **무해(harmless)**하기도 하길 바랍니다(폭탄 제조법을 알려주거나 잘못된 의학적 조언을 하는 등의 위험한 행동을 하지 않음).

보통 로봇을 안전하게 만들기 위해서는 처음부터 다시 설계를 하여 재학습(retraining)시켜야 합니다. 이는 규칙이 바뀔 때마다 학생을 다시 교육하기 위해 매번 새로운 선생님을 고용하는 것과 같습니다. 이는 비용이 많이 들고, 느리며, 많은 데이터가 필요합니다.

최근의 몇몇 방법들은 로봇을 재학습시키지 않고도, 로봇이 말하는 도중(추론 단계, inference)에 가이드하는 방법을 시도합니다. 하지만 이 방법들에는 결함이 있습니다. 보통 인간 운영자가 "마법의 숫자"를 직접 추측해야 한다는 점입니다.

예시: "좋아, 나쁜 단어 하나당 벌점을 5점으로 설정하자."
문제점: 만약 당신이 5를 선택했는데 로봇이 여전히 위험할 수 있습니다. 반대로 10을 선택하면, 로봇이 너무 겁을 먹어서 유용한 말을 전혀 하지 못하게 될 수도 있습니다. 이는 수학적 근거보다는 인간의 직관에 의존하는 "추측하고 확인하기(guess and check)" 게임입니다.

해결책: LARA (스마트한 "예산 관리" 시스템)

저자들은 LARA(Lagrangian Reward Augmentation)라고 불리는 새로운 프레임워크를 제안합니다. 마법의 숫자를 추측하는 대신, LARA는 유용함과 안전함 사이의 정확한 균형을 계산하기 위해 수학적 트릭을 사용합니다.

이것은 마치 엄격한 연료 예산이 있는 가족 자동차 여행과 같습니다.

1. 설정: 자동차와 지도

자동차: 얼어붙은(frozen) 언어 모델(로봇). 이미 만들어졌으며 우리는 엔진(가중치)을 변경하지 않습니다.
지도: "보상 모델(Reward Model)". 자동차가 목적지(유용함)에 도달하기 위해 얼마나 빨리 달릴 수 있는지 알려줍니다.
연료 게이지: "비용 모델(Cost Model)". 얼마나 많은 "위험"이나 "해로움"이 소모되고 있는지 측정합니다.
예산: 당신에게는 엄격한 "해로움" 한도가 있습니다 (예: "위험도는 10 유닛까지만 쓸 수 있음").

2. 기존 방식 vs LARA 방식

기존 방식 (수동 튜닝): 당신은 운전자에게 말합니다. "빨리 달리되, 연료를 너무 많이 쓰는 것 같으면 속도를 조금 줄여." 운전자는 얼마나 줄여야 할지 스스로 추측해야 합니다. 때때로 너무 빨리 달려 연료가 바닥나기도 하고(불안전), 때로는 너무 느리게 달려 목적지에 도착하지 못하기도 합니다(불유용).
LARA 방식 (Dual Variable 사용): LARA는 스마트한 GPS 계산기 역할을 합니다. 여행을 시작하기 전, 작은 샘플 경로(교정 세트, calibration set)를 살펴보고, 목적지에 가장 빠르게 도착하면서도 예산 내에서 움직일 수 있는 정확한 연료 가격을 계산합니다.
- LARA는 단 하나의 숫자(이를 $\lambda$ 또는 "람다"라고 부름)를 찾아냅니다.
- 이 숫자는 "안전의 그림자 가격(shadow price)"을 나타냅니다. 이것은 자동차에게 이렇게 말합니다: "위험을 1단위 감수할 때마다, 너는 유용함 X만큼을 잃게 된다."
- 그 후 자동차는 새로운 규칙에 따라 주행합니다: 유용함 - (람다 $\times$ 위험).

3. 작동 원리 ("교정" 단계)

LARA는 로봇 전체를 재학습시킬 필요가 없습니다. 그저 작은 "테스트 드라이브"(교정 세트)만 있으면 됩니다.

몇 가지 샘플 답변을 생성합니다.
이 답변들이 얼마나 "유용한지"와 얼마나 "위험한지"를 체크합니다.
빠른 수학적 탐색(온도 조절 장치의 완벽한 온도를 찾는 것과 유사함)을 실행하여, 유용함을 극대화하면서도 총 위험도를 한도 아래로 유지하는 특정 숫자( $\lambda$ )를 찾아냅니다.
이 숫자를 찾으면, 이를 로봇의 기존 "스코어링 시스템"에 대입합니다.

왜 특별한가?

이 논문은 LARA가 어떻게 작동하는지에 대해 두 가지 주요 사항을 주장합니다.

1. "Best-of-N" (리스트 중 최고의 답변 고르기)의 경우
로봇이 당신의 질문에 대해 20개의 서로 다른 답변을 작성한다고 가정해 봅시다.

LARA가 없다면: 당신은 가장 듣기 좋은 답변을 고를 수 있지만, 그것이 위험할 수도 있습니다.
LARA가 있다면: 시스템은 계산된 "람다"를 사용하여 20개의 답변 모두를 채점합니다. 이 방식은 안전 예산을 엄격히 준수하면서도 가장 유용한 답변을 자동으로 선택합니다. 논문은 이 방법이 완벽한 균형점을 찾는다는 것을 수학적으로 증명합니다.

2. "Token-Level" (단어 단위 가이드)의 경우
로봇이 문장을 한 단어씩 써 내려가는 상황을 상상해 보세요.

LARA가 있다면: 위험한 단어에 얼마만큼의 벌점을 줄지 추측하는 대신, 시스템은 계산된 "람다"를 사용하여 다음 단어가 나올 확률을 조정합니다. 이는 마치 정확한 제한 속도를 알고 있는 교통 경찰이 단순히 손을 흔들며 운이 좋기를 바라는 것이 아니라, 자동차가 제한 속도 바로 아래를 유지하도록 직접 지시하는 것과 같습니다. 논문은 이를 수학에 기반한 "원칙적인 휴리스틱(principled heuristic)", 즉 단순한 추측이 아닌 스마트한 규칙이라고 부릅니다.

결과

저자들은 두 가지 인기 있는 로봇 모델(Llama 및 Qwen)에 대해 테스트를 진행했습니다.

발견된 사실: LARA는 "추측" 기반의 다른 방법들에 비해 유용함과 안전함 사이의 균형을 훨씬 더 잘 맞추었습니다.
비교: LARA를 사용한 "Best-of-N" 방식은 막대한 비용이 드는 재학습 모델(수 주간의 훈련이 필요한 모델)과 거의 대등한 성능을 보였지만, 로봇의 뇌를 바꾸지 않고도 단 몇 초 만에 이를 수행했습니다.
트레이드오프: 로봇을 무용하게 만들지 않으면서도 해로운 행동을 성공적으로 차단했습니다.

요약 비유

AI를 안전하게 만드는 것이 집을 방화 구조로 만들기 위해 집을 새로 짓는 것이라면, LARA는 화재 크기에 따라 물의 압력을 자동으로 조절하는 스마트 스프링클러 시스템을 설치하는 것과 같습니다. 집을 새로 지을 필요 없이, 스프링클러를 한 번만 교정해 두면 집을 안전하게 유지하면서도 편안하게 생활할 수 있습니다.

기술 요약: 라그랑주 보상 증강(Lagrangian Reward Augmentation, LARA)을 통한 안전한 추론 시점 정렬

문제 정의

강화 학습으로부터의 인간 피드백(RLHF)은 대규모 언어 모델(LLM)을 인간의 선호도에 정렬하기 위한 표준 프레임워크입니다. 그러나 전통적인 정렬 방식은 정책 미세 조정(예: RLHF 또는 Direct Preference Optimization과 같은 DPO)에 의존하며, 이는 계산 비용이 많이 들고 안전 요구 사항이나 사용자 집단이 변경될 때마다 재학습이 필요합니다.

추론 시점 정렬(inference-time alignment) 방법들(Best-of-N 재순위화 또는 토큰 수준의 보상 유도 디코딩 등)은 가중치 업데이트 없이 보조 보상 신호를 사용하여 고정된 모델을 제어하는 방식으로 등장했지만, 현재는 **안전한 정렬(safe alignment)**을 위한 원칙적인 메커니즘이 부족합니다. 기존 방식들은 일반적으로 단일 스칼라 점수를 최적화합니다. 유용성(helpfulness)과 안전성(harmlessness) 목표가 충돌할 때, 실무자는 이 둘 사이의 트레이드오프를 수동으로 조정해야 합니다. 또한, 이러한 방법들은 디코딩 과정에서 명시적인 안전 제약 조건(예: 기대 비용에 대한 예산 제한)을 본질적으로 강제하지 못하므로, 안전하지 않은 출력이 생성되거나 불필요하게 유용성이 감소하는 문제가 발생합니다.

방법론: 라그랑주 보상 증강 (LARA)

저자들은 Safe RLHF의 제약 최적화 구조를 추론 단계로 전이시키는 디코더 불가지론적(decoder-agnostic) 프레임워크인 LARA를 제안합니다. LARA는 모델 가중치를 업데이트하는 대신, 스칼라 쌍대 변수(dual variable)를 보정하여 증강된 보상 신호를 생성합니다.

1. 이론적 정식화

이 방법은 참조 정책( $\pi_{ref}$ )에 대한 KL-발산(KL-divergence) 항으로 정규화하면서, 기대 비용 제약( $c \le \tau$ )을 만족하도록 하는 Safe RLHF 목적 함수에서 시작합니다:
$\max_{\pi} \mathbb{E}[r(x,y)] - \beta D_{KL}(\pi \parallel \pi_{ref}) \quad \text{s.t.} \quad \mathbb{E}[c(x,y)] \le \tau$

**라그랑주 쌍대성(Lagrangian duality)**을 사용하여, 저자들은 안전 제약 조건을 쌍대화합니다. 이를 통해 무한 차원의 정책 최적화 문제를 비음수 쌍대 변수 $\lambda$ (라그랑주 승수)에 대한 일차원 볼록 최적화 문제로 변환합니다.

증강된 보상: 쌍대 변수는 증강된 보상 $r_\lambda(x, y) = r(x, y) - \lambda c(x, y)$ 를 정의합니다.
쌍대 목적 함수: 쌍대 함수 $g(\lambda)$ 는 볼록함(convex)임이 입증되었습니다. 그 도함수는 $r_\lambda$ 에 의해 유도된 Gibbs-tilted 정책 하에서의 안전 예산 $\tau$ 와 기대 비용 간의 차이에 해당합니다.
최적성: 최적의 $\lambda^*$ 는 상보 여유 조건(complementary slackness condition)을 만족하여, (제약 조건이 활성화된 경우) 기대 비용이 예산 $\tau$ 와 같도록 보장하는 동시에 유용성을 최대화합니다.

2. 경험적 보정 (Empirical Calibration)

LARA는 최적해를 분석적으로 풀 필요가 없습니다. 대신 소규모 데이터셋을 통해 $\lambda^*$ 를 경험적으로 보정합니다:

데이터: 소수의 프롬프트를 샘플링하고, 고정된 참조 정책으로부터 여러 응답을 생성합니다.
추정: 컴팩트한 구간 $[0, \Lambda]$ 내에서 경험적 쌍대 목적 함수를 최소화함으로써 쌍대 변수를 추정합니다.
알고리즘: 쌍대 함수의 볼록성과 단조성 덕분에, **이진 탐색(binary search)**을 사용하여 최적의 $\lambda$ 를 효율적으로 찾을 수 있습니다.
강건성: 보정과 배포 사이의 분포 불일치를 처리하기 위해, 저자들은 보정 세트를 부트스트래핑하고 추정된 $\lambda$ 분포의 상위 분위수를 보수적인 선택지로 사용하는 것을 제안합니다.

3. 디코딩과의 통합

보정이 완료되면, 스칼라 $\hat{\lambda}$ 는 증강된 보상 $r_{\hat{\lambda}}$ 를 계산하는 데 사용됩니다. 이 점수는 기존의 추론 시점 정렬 방법으로 전달되는 "드롭인(drop-in)" 신호 역할을 합니다:

시퀀스 수준 (예: Best-of-N): 보정된 $\lambda$ 는 기대 비용 제약 문제를 해결하는 정확한 해를 제공합니다.
토큰 수준 (예: 보상 유도 탐색): 보정된 $\lambda$ 는 제약된 목적 함수로부터 유도된 원칙적인 휴리스틱을 제공하지만, 그리디(greedy)한 특성의 토큰 수준 디코딩 때문에 인구 집단 수준의 제약 조건을 정확히 만족시키지는 못할 수 있습니다.

주요 기여

정식화: 본 논문은 안전한 추론 시점 정렬을 Safe RLHF 목적 함수의 쌍대로 정식화하여, 문제를 일차원 볼록 보정 작업으로 축소했습니다.
기하학 및 효율성: 저자들은 쌍대 목적 함수의 기하학적 구조를 규명하여 그 볼록성과 도함수의 단조성을 증명했습니다. 이를 통해 효율적인 이진 탐색 보정 절차를 가능하게 했습니다.
이론적 보장: 추정량 $\hat{\lambda}_{N,K}$ 에 대한 유한 표본 보장이 제공되며, 오차가 프롬프트 수( $N$ )와 프롬프트당 샘플 수( $K$ )에 의해 제어되면서 모집단 최소화 도 $\lambda^*$ 로 집중됨을 보여줍니다.
디코더 불가지론적 프레임워크: LARA는 기존의 추론 시점 정렬 방법과 호환되는 보정 레이어로 설계되었으며, 수동적인 스칼라 튜닝을 원칙적인 안전 제약 기반 구축으로 대체합니다.

실험 결과

저자들은 BeaverTails 데이터셋을 사용하여 Llama-3.2-3B 및 Qwen-2.5-3B 모델에 대해 LARA를 평가했습니다.

제약 조건 만족: 실험을 통해 보정된 $\lambda^*$ 가 안전 제약의 경계에 위치함을 확인했습니다. $\lambda < \lambda^*$ 인 경우 안전 예산을 위반했고, $\lambda > \lambda^*$ 인 경우 제약 조건은 만족했으나 유용성이 불필요하게 감소했습니다.
베이스라인 대비 성능:
- 추론 시점 방법: LARA와 결합된 **Best-of-N (BoN)**이 가장 우수한 유용성-무해성 트레이드오프를 달ian했습니다. BoN은 가중치 업데이트 없이도 SafeDPO보다 유용성 측면에서 뛰어난 성능을 보였으며 무해성 측면에서도 경쟁력 있는 결과를 보였습니다.
- 토큰 수준 방법: 토크 수준 방법들(ARGS, PARGS, CD-Q)은 일반적으로 BoN에 비해 유용성 측면에서 낮은 성능을 보였으나, CD-Q는 $\lambda$ 가 증가함에 따라 강력한 무해성 개선을 보여주었습니다.
- 미세 조정 대비: LARA 기반의 BoN은 미세 조정 기반 베이스라인(Safe-RLHF, SafeDPO)의 성능에 근접했습니다. 다만, 미세 조정된 방법들이 전반적인 트레이드오프에서는 여전히 가장 우수한 성과를 냈습니다. 그러나 LARA는 배포 유연성과 비용 측면에서 큰 이점을 제공합니다.
평가: 결과는 모델 기반 평가(BeaverTails 보상/비용 모델 사용)와 GPT-4o-mini 기반 평가 모두에서 일관되게 나타났으나, GPT는 보상 모델에 비해 $\lambda$ 변화에 대한 민감도가 낮았습니다.

의의 및 주장

본 논문은 LARA가 안전한 추론 시점 정렬을 위한 수동적인 페널티 튜닝의 원칙적인 대안을 제공한다고 주장합니다. 안전 제약 조건을 쌍대화함으로써, LARA는 모델 가중치를 수정하지 않고도 안전 예산을 디코딩 목적 함수로 변환합니다.

시퀀스 수준 샘플러의 경우: LARA는 기대 비용에 대한 정확한 해석을 제공하며, 이론적으로 제약된 목적 함수 하에서의 최적의 트레이드오프를 보장합니다.
토큰 수준 디코더의 경우: 비록 인구 집단 제약 조건을 엄격하게 강제할 수는 없더라도, 제약된 목적 함수의 기하학적 구조를 존중하는 원칙적인 휴리스틱을 제공합니다.

저자들은 LARA가 안전 제약 조건을 훈련 시점의 정렬에서 실제 디코딩 시점의 제어로 성공적으로 전이시켜, 유용성과 무해성 사이의 트레이드오프를 개선하고 추론 시점 방법들을 미세 조정 기반 접근 방식의 성능에 가깝게 끌어올렸다고 결론지었습니다. 이 연구는 반복적인 재학습의 계산 오버헤드를 피하면서도 안전 제약 조건을 추론 시점에 효과적으로 관리할 수 있음을 강조합니다.