🤖 machine learning

Amortizing intractable inference in large language models

이 논문은 대규모 언어 모델에서 난해한 추론을 분할 계산(amortize)하기 위해 GFlowNet을 사용하는 방안을 제안하며, 이를 통해 전통적인 최대 가능도 또는 보상 극대화 학습의 데이터 효율적인 대안으로서 제약 조건이 있는 생성 및 사고 사슬(chain-of-thought) 추론과 같은 작업에 대해 복잡한 사후 분포로부터 샘플링하는 것을 가능하게 한다.

원저자: Edward J. Hu, Moksh Jain, Eric Elmoznino, Younesse Kaddar, Guillaume Lajoie, Yoshua Bengio, Esmeralda S. Whitammer

게시일 2026-09-01

📖 3 분 읽기☕ 가벼운 읽기

CC BY 4.0

원저자: Edward J. Hu, Moksh Jain, Eric Elmoznino, Younesse Kaddar, Guillaume Lajoie, Yoshua Bengio, Esmeralda S. Whitammer

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

대규모 언어 모델은 인간 지식의 방대한 도서관이며, 앞선 단어를 바탕으로 문장의 다음 단어를 예측하는 디지털 네트워크로 압축되어 있습니다. 이 모델들은 시작점이 주어졌을 때 이야기를 이어가거나 질문에 답하는 데 탁로하도록 훈련되었으며, 이 과정은 왼쪽에서 오른쪽으로 자연스럽게 흐릅니다. 그러나 우리가 이 모델들에게 수행시키고자 하는 가장 흥미로운 과제 중 상당수는 역방향으로 생각하거나 이야기의 중간 부분을 채우는 것을 요구합니다. 이야기의 시작과 끝을 주고 그 중간을 만들어내라고 요청받거나, 특정 답변 뒤에 숨겨진 추론 과정을 설명하라고 요청받는 상황을 상상해 보십시오. 이러한 시나리오에서 모델은 단순히 다음 단어를 추측하는 것에 그칠 수 없습니다. 모델은 시작과 끝을 연결하는 특정한 경로를 찾기 위해 거대하고 복용한 가능성의 풍경을 탐색해야 합니다. 이는 가능한 경로의 수가 너무 많아 하나씩 확인하는 것이 불가능하기 때문에 매우 어려운 수학적 문제입니다. 또한, 이 풍경을 탐색하는 표준적인 방법들은 종종 한두 개의 흔한 경로에만 머물러 존재할 수 있는 풍부하고 다양한 유효한 해결책들을 놓치곤 합니다.

Mila – 퀘벡 AI 연구소와 옥스퍼드 대학교의 연구진은 이러한 모델들이 그 복잡한 풍경을 탐색하는 법을 가르치는 새로운 방법을 개발했습니다. 모델이 단순히 '최선의' 답변에 대해 점수를 극대화하도록 훈련하는 대신(이는 종종 반복적이고 좁은 사고로 이어집니다), 그들은 '아모티제이션된 추론(amortized inference)'이라 불리는 방법을 사용했습니다. 이 접근 방식은 문제를 단 하나의 완벽한 답을 찾는 것이 아니라, 다양한 정답 집합을 찾는 과정으로 취급합니다. 이를 위해 그들은 '생성적 흐름 네트워크(generative flow network)'라고 알려진 기술을 도입했는데, 이는 모델이 하나의 과도하게 사용된 패턴으로 붕괴되지 않고 다양한 유효한 추론 사슬을 탐색할 수 있도록 전체 범위의 솔루션으로부터 샘플링하는 법을 배우는 가이드 역할을 합니다.

연구팀은 산술 계산이나 영화 리뷰가 의견을 나타내는지 사실을 나타내는지 분류하는 것과 같이 다단계 추론이 필요한 문제를 풀도록 대규모 언어 모델을 미세 조정함으로써 이를 입증했습니다. 한 실험에서, 그들은 모델에게 이야기의 시작과 끝이 주어졌을 때 짧은 이야기의 누락된 중간 문장을 채우도록 요청했습니다. 표준적인 방법들은 문법적으로는 맞지만 서사의 흐름에는 맞지 않는 문장을 생성하는 경우가 많았습니다. 그러나 새로운 방법은 시작과 끝을 일관되게 연결하는 중간 문장을 성공적으로 생성하여, 전체 맥락을 이해하는 능력이 훨씬 높음을 보여주었습니다. 계산기를 도구로 사용하는 간단한 수학 문제를 포함한 또 다른 테스트에서는, 모델에게 계산 과정을 단계별로 나누도록 요청했습니다. 기존의 훈련 방식들은 모델이 숫자를 반복하거나 도구를 올바르게 사용하지 못하게 만든 반면, 새로운 접근 방식은 모델이 단계를 신중하게 계획하도록 가르쳐 특히 처음 보는 문제들에 대해 정확도를 극적으로 향상시켰습니다.

결과는 이 방법이 데이터가 부족할 때 특히 강력하다는 것을 시사합니다. 단 10개의 영화 리뷰 예시만을 학습한 테스트에서, 이 새로운 방법은 표준 훈련 기술보다 현저히 높은 정확도를 달aba 성취했으며, 50개의 예시가 있을 때도 계속해서 이를 능가했습니다. 연구진은 모델이 다양한 추론 경로를 샘플링하도록 장려하고 그 결론들을 결합함으로써 시스템이 더 견고하고 신뢰할 수 있게 된다는 것을 발견했습니다. 이는 모델이 단순히 문제를 해결하는 단 하나의 방법을 암기하는 것이 아니라, 그것을 어떻게 추론해 나가는지에 대한 근본적인 구조를 배우고 있다는 점에서 매우 중요한 차이입니다. 이 연구는 목표를 '가장 가능성 높은 단 하나의 답 찾기'에서 '정답의 전체 다양성 탐색'으로 전환함으로써, 우리가 이 강력한 도구들을 더 유연하게 만들고 복잡한 추론 과제에 더 적합하게 만들 수 있음을 보여줍니다.

또한 이 작업은 현재 모델들이 훈련되는 방식의 한계를 강조합니다. 모델이 보상을 극대화하도록 강요될 때, 모델은 높은 점수를 얻기 위해 동일한 구절을 반복하거나 실제 논리를 무시하는 등의 지름길을 찾는 경향이 있습니다. 연구진은 그들의 방법이 가능한 솔루션의 전체 분포를 일치시킴으로써, 모델이 단일하고 결함이 있는 사고 패턴으로 붕괴되지 않도록 하여 이러한 함정을 피한다는 것을 보여주었습니다. 이 접근 방식은 모델이 훈련받은 것보다 더 많은 숫자가 포함된 산술 문제를 푸는 것과 같이 새로운 상황에 더 잘 일반화할 수 있게 해줍니다. 이러한 발견은 인공지능을 단순한 패턴 매칭을 넘어, 논거를 구성하고 단계별로 문제를 해결하는 법을 이해하는 더 미묘한 수준의 진정한 추론을 할 수 있는 방향으로 나아가게 하는 유망한 길을 제시합니다.

기술 요약: 대규모 언어 모델에서의 난해한 추론에 대한 분할 상환(Amortizing)

문제 정의

자기회귀적 대규모 언어 모델(LLM)은 지식을 다음 토큰 조건부 분포로 압축하며, 이로 인해 트랙터블(tractable)한 추론은 접두사(prefix)에 조건화된 시작-끝 샘플링으로 제한됩니다. 그러나 문장 이어 쓰기(sequence continuation), 텍스트 인필링(text infilling), 제약 조건이 있는 생성(constrained generation), 그리고 사고의 사슬(chain-of-thought) 추론과 같은 많은 중요한 작업들은 **난해한 사후 분포(intractable posterior distributions)**로부터 샘플링하는 것을 요구합니다.

구체적으로, 인필링(시작 $X$ 와 끝 $Y$ 가 주어졌을 때 중간 세그먼트 $Z$ 를 생성하는 것)이나 추론(프롬프트 $X$ 가 주어졌을 때 결과 $Y$ 로 이어지는 잠재적 사고의 사슬 $Z$ 를 찾는 것)과 같은 작업은 사후 분포 $p(Z|X, Y)$ 를 계산하는 문제를 포함합니다. 결합 가능도(joint likelihood) $p_{LM}(XZY)$ 는 계산하기 쉽지만, $Z$ 를 사후 분포로부터 토큰 단위로 샘플링하는 데 필요한 조건부 분포는 계산이 불가능(intractable)합니다.

기존 접근 방식은 다음과 같은 상당한 한계에 직면해 있습니다:

마르코프 연쇄 몬테카를로(MCMC): 다중 모드(multi-modal) 언어 데이터에 적합한 제안 분포(proposal distribution)를 설계하기 어렵고, 새로운 입력에 대해 지나치게 느린 경우가 많습니다.
강화 학습(RL) / 정책 최적화(예: PPO): 이러한 방법들은 보상을 최대화하지만, 타겟 분포의 전체 다양성을 포착하지 못하고 소수의 모드로 붕괴(collapse)되는 경st가 있습니다. 이는 보상 함수가 잘못 지정되었을 때 "과최적화된 샘플러(overoptimized samplers)"를 초래합니다.
지도 미세 조정(SFT): 타겟 분포로부터의 샘플에 대한 접근이 필요하지만, 난해한 문제에서는 이러한 샘플을 일반적으로 얻을 수 없습니다.

방법론

저자들은 이러한 난해한 사후 분포로부터 샘플링하도록 LLM을 훈련시키기 위해 **생성적 플로 네트워크(GFlowNets)**를 통한 **분할 상환 베이지안 추론(amortized Bayesian inference)**을 사용할 것을 제안합니다.

핵심 프레임워크

이 접근 방식은 입력 $X$ 가 잠재 변수 $Z$ (예: 추론 사슬 또는 인필링된 텍스트)를 통해 출력 $Y$ 를 생성하는 잠재 변수 모델로 문제를 취급합니다. 목표는 $p_{LM}(Z|X, [Y]) \propto p_{LM}(XZY)$ 로부터 샘플링하는 정책 $q_{GFN}(Z|X, [Y])$ 를 훈련시키는 것입니다.

GFlowNet 목적 함수: 이 방법은 다양성을 추구하는 강화 학습 알고리즘의 일종인 GFlowNet을 활용합니다. 기대 보상을 최대화하는 표준 RL과 달리, GFlowNet은 주어진 보상 함수 $R(Z)$ 에 비례하는 객체를 샘플링하도록 정책을 훈련합니다.
보상 정의: 보상은 기본 LLM이 할당한 결합 가능도로 정의됩니다: $R(Z) = p_{LM}(XZY)$ . 이는 효과적으로 타겟 분포를 베이지안 사후 분포로 설정합니다.
훈련 절차:
- 초기화: GFlowNet 정책은 사전 훈련된 LLM의 복사본으로 초기화됩니다(효율성을 위해 종종 LoRA로 미세 조정됨).
- 학습 목적 함수: 저자들은 수정된 하위 궤적 균형(Subtrajectory Balance, SubTB) 목적 함수를 사용합니다. 이 목적 함수는 궤적이 임의의 상태에서 종료될 수 있는 경우에도 완전한 시퀀스를 생성할 확률이 보상에 비례하도록 보장합니다.
- 분할 상환(Amortization): 정책을 입력 $X$ (및 선택적으로 $Y$ )에 조건화함으로써, 모델은 일반화 가능한 샘플러를 학습합니다. 이를 통해 각 쿼리에 대해 비용이 많이 드는 MCMC와 같은 샘플링 절차를 다시 실행하지 않고도 새로운 입력에 대해 효율적인 추론이 가능합니다.
- 탐색(Exploration): 조합론적 탐색 공간을 처리하기 위해, 훈련 배치(batch)는 현재 정책, 템퍼링된(tempered) 버전의 정책, 그리고 과거 궤적의 리플레이 버퍼에서 추출한 궤적들로 구성됩니다.

사고의 사슬(Chain-of-Thought)에 대한 적용

본 논문은 사고의 사슬 추론을 잠재 변수 모델링 문제로 해석합니다. 모델에게 단일 사슬을 생성하도록 프롬프트를 주는 대신, GFlowNet으로 미세 조정된 모델은 다양한 잠재적 사슬 $Z$ 를 샘플링합니다. 최종 예측은 **베이지안 모델 평균화(Bayesian Model Averaging)**를 통해 얻습니다. 즉, GFlowNet에서 여러 $Z$ 를 샘플링하고, 그로부터 파생된 기본 LLM의 $Y$ 예측값들을 집계합니다.

주요 기여

일반 알고리즘: GFlowNet 목적 함수를 사용하여 난해한 LLM 사후 분포로부터 분할 상환 샘플링을 수행하는 일반적인 방법을 제시합니다.
확률적 추론 접근법: 사고의 사슬 추론을 베이지안 추론 문제로 공식화하여, 다단계 합리화 작업에 대한 데이터 효율적인 적응을 가능하게 합니다.
실증적 검증: 다음 네 가지 서로 다른 도메인에서 방법론을 입증했습니다:
- 문장 이어 쓰기 (가능도를 유지하면서 다양성 극대화).
- 스토리 인필링 (일관된 중간 세그먼트 생성).
- 주관성 분류 (잠재적 근거를 활용한 퓨샷 학습).
- 도구 사용을 포함한 정수 산술 (단계별 추론).

실험 결과

저자들은 표준 샘플링, 다양한 빔 서치(diverse beam search), 지도 미세 조정(SFT), 그리고 근사 정책 최적화(PPO)를 포함한 베이스라인들과 비교 평가를 진행했습니다.

문장 이어 쓰기: GFlowNet 미세 조정은 베이스라인이 5배의 연산량을 제공했을 때조차 다양한 빔 서치보다 높은 로그 가능도를 유지하면서 더 높은 샘플 다양성을 달ach했습니다.
스토리 인필링: ROCStories 데이터셋에서, GFlowNet 미세 조정은 BERTScore, BLEU-4, GPT-4 평가에서 프롬프팅 및 SFT를 유의미하게 능가하며, 이야기의 시작과 끝을 더 잘 연결하는 인필을 생성했습니다.
주관성 분류 (퓨샷): 6B 파라미터 모델을 사용하여 단 10개의 레이블된 예시만으로, GFlowNet 미세 조정은 **71.4%**의 정확도를 달성하여 SFT(64.3%)와 퓨샷 프롬프팅(61. 피드백)을 능가했습니다.
도구 사용을 포함한 정수 산술: 5개의 피연산자가 포함된 분포 외(OOD) 작업(3-4개로 학습)에서, GFlowNet 미세 조정은 **40.7%**의 정확도를 달성했습니다. 이는 SFT(12.8%)와 PPO(5.6%)를 크게 앞선 수치입니다. 저자들은 PPO가 보상 오설정(reward misspecification)으로 인해 계산기 도구를 유효하게 호출하지 않는 가짜 시퀀스를 생성하는 문제를 겪은 반면, GFlowNet의 분포 매칭 접근 방식은 이에 견고했음을 언급했습니다.

의의 및 주장

본 논문은 GFlowNet 미세 조정이 최대 가능도 학습 및 보상 최대화 정책 최적화에 대한 원칙적인 대안을 제공한다고 주장합니다. 이 방법의 주요 의의는 다음과 같습니다:

다양성 및 견고성: 단일 보상을 최대화하는 대신 전체 타겟 분포를 매칭함으로써, 모드 붕괴를 방지하고 보상 오설정에 대해 견고합니다.
데이터 효율성: 분할 상환 추론 접근 방식은 매우 적은 양의 레이블된 예시만으로도 복잡한 추론이 필요한 작업에 효과적으로 적응할 수 있게 합니다.
일반화: 표준 RL 방법이 솔루션 공간을 효과적으로 탐색하지 못하는 추론 작업에서 우수한 분포 외(OOD) 일반화 성능을 보여줍니다.

저자들은 이 패러다임이 추가적인 데이터 없이도 계산량을 더 나은 테스트 타임 성능으로 전환하여, LLM의 더 유연한 추론을 가능하게 한다고 결론짓습니다. 또한, 사용된 모델의 규모(최대 6B 파라미터)와 이 방법이 추론 능력을 향상시키기는 하지만, 기본 모델의 지식 표현에 뿌리를 둔 환각(hallucination)이나 미교정(miscalibration) 문제를 본질적으로 해결하는 것은 아니라는 점을 인정합니다.