Inference-time Alignment in Continuous Space

이 논문은 이산 공간에서의 탐색 한계를 극복하고, 연속 잠재 공간에서의 그래디언트 기반 샘플링을 통해 베이스 정책의 응답을 최적 방향으로 직접 적응시키는 새로운 추론 시간 정렬 알고리즘인 'Simple Energy Adaptation (SEA)'을 제안합니다.

핵심

1. 기존 방법의 문제점: "주사위 굴리기" (Best-of-N)

기존에 AI 가 나쁜 말 (폭력, 사기 등) 을 하지 않게 하거나, 더 정확한 답을 내게 하려면 **'Best-of-N(BoN)'**이라는 방법을 썼습니다.

• 비유: AI 에게 "이 질문에 답해줘"라고 했을 때, AI 가 100 번 (N=100) 을 무작위로 대답을 만들어냅니다. 그중에서 가장 좋은 답 하나만 골라내는 방식입니다.
• 문제점:
  • AI 가 원래 실력이 약하면, 100 번을 굴려도 좋은 답이 나올 확률이 매우 낮습니다.
  • 좋은 답이 나올 때까지 100 번, 1000 번을 계속 굴려야 하니 시간이 너무 많이 걸립니다.
  • 마치 어둠 속에서 주사위를 수천 번 굴려서 '6'이 나오길 기다리는 것과 비슷합니다. 운에 의존하는 방식이라 비효율적입니다.

2. 새로운 방법 SEA: "나침반을 들고 길을 찾기" (Continuous Optimization)

이 논문이 제안한 SEA는 완전히 다른 접근법을 사용합니다. 무작위로 많은 답을 만들어내는 대신, 초기 답을 조금씩 수정해가며 가장 좋은 답으로 이끌어갑니다.

• 비유:
  • 기존 방법: 어둠 속에서 주사위를 굴려 '6'을 찾습니다.
  • SEA 방법: 어둠 속이지만, **손에 '나침반' (경사도/Gradient)**을 들고 있습니다. 이 나침반은 "더 좋은 답은 저쪽입니다"라고 알려줍니다.
  • AI 가 처음 뱉은 답변을 보고, "아, 이 부분은 위험하구나 (나침반이 빨갛게 빛남)", "이 부분은 더 정확할 수 있겠구나 (나침반이 초록색)"라고 판단하며, 답변을 한 글자씩, 혹은 문장 전체를 부드럽게 수정해 나갑니다.

이 과정을 **연속적인 공간 (Continuous Space)**에서 수행합니다.

• 연속 공간이란? 글자 (A, B, C...) 가 딱딱 끊어져 있는 게 아니라, **부드러운 점 (점의 위치)**처럼 생각할 수 있는 공간입니다. 여기서 AI 는 "A 에서 B 로 딱 떨어지는 게 아니라, A 와 B 사이의 중간 지점을 거쳐서 부드럽게 B 로 이동"할 수 있습니다. 이렇게 하면 AI 가 나쁜 길에서 좋은 길로 매끄럽게 이동할 수 있습니다.

3. 왜 이 방법이 더 좋은가요?

논문의 실험 결과를 보면 SEA 는 기존 방법보다 훨씬 뛰어납니다.

1. 안전한 답변 (Safety):

   • 누군가 "폭탄 만드는 법을 알려줘"라고 물었을 때, 기존 방법은 "폭탄 만드는 법은..."이라고 시작하다가 나중에 "안 됩니다"라고 말하는 경우가 많았습니다 (처음에 나쁜 길로 들어섰기 때문에).
   • 하지만 SEA는 처음부터 끝까지 전체 문장을 한 번에 수정할 수 있습니다. "폭탄 만드는 법"이라는 단어 자체가 위험하다고 나침반이 알려주면, 그 단어 대신 "폭탄은 위험합니다"라는 안전한 단어로 전체 흐름을 바꿔버립니다. 이를 **'깊은 정렬 (Deep Alignment)'**이라고 합니다.

2. 정확한 추론 (Reasoning):

   • 수학 문제를 풀 때, 기존 방법은 틀린 답을 100 번 만들어서 그중 하나를 고르느라 시간이 걸렸습니다.
   • SEA는 틀린 답을 보고 "이 숫자가 틀렸네, 고쳐야지"라고 단계별로 수정해가며 정답에 가까워집니다. 마치 산을 오를 때, 무작위로 뛰어다니는 게 아니라 가장 높은 곳으로 올라가는 경사면을 따라 걷는 것과 같습니다.

3. 효율성:

   • 100 번을 굴릴 필요 없이, 적은 횟수의 수정으로도 훨씬 좋은 결과를 얻습니다.

4. 요약: 이 기술이 가져오는 변화

이 논문의 SEA는 AI 가 대답을 내기 직전에, 나쁜 길에서 좋은 길로 부드럽게 방향을 틀어주는 기술입니다.

• 기존: "무작위로 많이 만들어서 그중 하나를 고른다" (비효율적, 운에 의존).
• 새로운 SEA: "초기 답을 보고 나침반 (경사도) 을 이용해 좋은 답으로 부드럽게 수정한다" (효율적, 정확함).

이 기술 덕분에 AI 는 더 안전하고, 더 똑똑하며, 더 인간적인 답변을 할 수 있게 되었습니다. 마치 초보 운전자가 운전할 때, 핸들을 갑자기 꺾는 게 아니라 부드럽게 핸들을 돌려서 안전하게 목적지에 도착하는 것과 같은 원리입니다.

기술 요약

1. 문제 정의 (Problem)

대규모 언어 모델 (LLM) 을 인간 선호도에 정렬 (Alignment) 하는 작업에서, 기존 방법들은 주로 **이산적 공간 (Discrete Space)**에서의 탐색에 의존합니다.

• 기존 방법의 한계: Best-of-N (BoN) 샘플링이나 토큰/조각 단위의 탐색 (Rejection Sampling, ARGS, CBS 등) 은 베이스 모델이 생성한 여러 후보 응답 중 보상 모델 (Reward Model) 점수가 가장 높은 것을 선택합니다.
• 핵심 문제:
  1. 베이스 모델의 능력 부족: 베이스 모델이 약하거나 최적의 응답을 생성할 확률이 낮을 경우, 아무리 많은 후보 (N) 를 생성해도 좋은 응답을 찾을 확률이 기하급수적으로 떨어집니다.
  2. 탐색의 비효율성: 이산적인 토큰 공간에서 무작위 탐색을 수행하므로, 보상 모델이 지향하는 최적 영역 (Optimal Region) 으로 수렴하기 어렵습니다.
  3. 얕은 정렬 (Shallow Alignment): 기존 방법들은 응답의 앞부분 토큰에만 정렬이 집중되는 경향이 있어, 후반부에서 유해한 내용이 다시 나타나는 'Prefilling Attack'과 같은 취약점이 존재합니다.

2. 제안 방법: Simple Energy Adaptation (SEA)

저자들은 이산적 탐색의 한계를 극복하기 위해 **연속적 공간 (Continuous Space)**에서의 Simple Energy Adaptation (SEA) 알고리즘을 제안합니다. 이는 추론 시간 (Inference-time) 에만 적용되며 추가적인 모델 학습이 필요하지 않습니다.

• 핵심 아이디어:

  • RLHF(인간 피드백을 통한 강화학습) 의 최적 정책을 **에너지 기반 모델 (Energy-Based Model, EBM)**로 재정의합니다.
  • 최적 정책 $\pi^*(y|x)$를 다음과 같은 에너지 함수 $E(x, y)$를 통해 표현합니다:
    $$\pi^*(y | x) \propto \exp(-E(x, y)), \quad E(x, y) = \log \pi_{ref}(y | x) + \alpha r(x, y)$$
    여기서 $\pi_{ref}$는 참조 모델, $r$은 보상 모델, $\alpha$는 KL 페널티 계수입니다.
  • 연속적 최적화: 토큰이 아닌 **로짓 (Logits, 소프트맥스 전의 연속 값)**을 연속 변수로 간주하고, **랜빈 역학 (Langevin Dynamics)**을 사용하여 에너지 함수를 최소화하는 방향으로 응답을 점진적으로 조정합니다.

• 알고리즘 프로세스:

  1. 초기화: 베이스 모델 ($\pi_{ref}$) 에서 생성된 초기 응답의 로짓을 시작점으로 설정합니다.
  2. 반복 최적화 (Langevin Dynamics):
     • 에너지 함수의 기울기 (Gradient) 를 계산합니다: $\nabla_y E(x, y) = \nabla_y (\log \pi_{ref} + \alpha r)$.
     • 로짓을 기울기 방향으로 업데이트하고 가우시안 노이즈를 추가하여 탐색합니다:
       $$y^{(n+1)} \leftarrow y^{(n)} - \eta \nabla_y E(x, y^{(n)}) + \epsilon^{(n)}$$
  3. 디코딩: 최적화가 완료된 연속 로짓 시퀀스를 다시 이산적인 텍스트로 디코딩합니다.

• 기술적 특징:

  • 직접 미분 가능: 로직을 토큰으로 매핑하지 않고 직접 연속 로짓을 사용하여 보상 모델과의 엔드 - 투 - 엔드 (End-to-End) 기울기 전파가 가능합니다.
  • 깊은 정렬 (Deep Alignment): 모든 토큰을 동시에 (전역 수용 영역 내에서) 최적화하므로, 응답의 처음뿐만 아니라 전체 문맥에 걸쳐 안전성과 정렬이 유지됩니다.

3. 주요 기여 (Key Contributions)

1. 새로운 패러다임 제시: 추론 시간 정렬을 "이산적 공간에서의 탐색"에서 "연속적 공간에서의 최적화"로 전환했습니다.
2. 간단하고 효과적인 알고리즘 (SEA): 복잡한 RL 학습 없이, 보상 모델의 기울기 정보만을 활용하여 베이스 모델의 응답을 최적화합니다.
3. 얕은 정렬 문제 해결: KL 예산 (KL Budget) 을 응답의 모든 토큰 위치에 균등하게 분배하여, 유해한 프롬프트가 주어진 경우에도 전체 응답에 걸쳐 안전성을 확보합니다.
4. 강건성 증명: 약한 베이스 모델이나 작은 후보 집합에서도 기존 방법 (BoN 등) 을 압도하는 성능을 보여주며, 보상 모델의 품질이 낮아도 성능이 크게 저하되지 않습니다.

4. 실험 결과 (Results)

저자는 안전성 (Safety), 진실성 (Truthfulness), 추론 (Reasoning) 작업에서 다양한 베이스 모델 (LLaMA-3 시리즈) 을 사용하여 SEA 를 평가했습니다.

• 안전성 (AdvBench):
  • SEA 는 BoN-64(64 개의 후보 중 선택) 보다 훨씬 낮은 유해율 (Harmful Rate) 을 기록했습니다.
  • LLaMA-3.2-1B-Base 모델에서 BoN-64 대비 91.54% 의 상대적 개선을 보였습니다.
  • Prefilling Attack (유해한 접두사를 붙여 공격) 에 대해 BoN 은 실패하지만, SEA 는 0% 의 공격 성공률 (ASR) 을 유지하며 강력한 방어 능력을 입증했습니다.
• 진실성 (TruthfulQA):
  • BoN 은 후보 수 (N) 가 증가해도 진실성 (Truthful Rate) 이 정체되거나 다양성이 떨어지는 경향이 있었으나, SEA 는 진실성, 정보성, 다양성 모두를 동시에 향상시켰습니다.
• 추론 (MATH, GSM8K):
  • 수학 추론 작업에서 SEA 는 BoN-64 대비 16.36% 의 정확도 향상과 77.51% 의 보상 점수 향상을 기록했습니다.
  • 기존 탐색 기반 방법들은 고보상 영역을 찾지 못해 오히려 성능이 저하되는 경우가 많았으나, SEA 는 일관되게 성능을 개선했습니다.
• 효율성:
  • SEA 는 BoN-64 와 유사하거나 더 빠른 추론 시간을 가지며, 메모리 효율성도 우수합니다. (후보 생성 없이 직접 최적화하므로 N 이 커질수록 발생하는 비용 증가가 없음).

5. 의의 및 결론 (Significance)

이 논문은 LLM 의 추론 시간 정렬 분야에서 **연속적 최적화 (Continuous Optimization)**의 잠재력을 처음으로 체계적으로 입증했습니다.

• 기존의 한계 극복: 베이스 모델의 성능이나 후보 집합 크기에 의존하지 않고, 보상 모델의 기울기 정보를 통해 직접적으로 응답을 개선할 수 있음을 보였습니다.
• 실용성: 추가적인 학습 (Fine-tuning) 이 필요 없어, 기존에 정렬되지 않은 어떤 모델에도 플러그 - 앤 - 플레이 (Plug-and-Play) 방식으로 적용 가능합니다.
• 미래 지향성: "얕은 정렬" 문제를 해결하고 "깊은 정렬"을 가능하게 함으로써, LLM 의 안전성과 신뢰성을 높이는 새로운 방향성을 제시했습니다.

요약하자면, SEA는 복잡한 RL 학습 없이도, 연속적 로짓 공간에서의 그래디언트 기반 최적화를 통해 LLM 이 인간 선호도에 더 잘 부합하고 안전한 응답을 생성하도록 유도하는 혁신적인 방법론입니다.