Thompson Sampling via Fine-Tuning of LLMs

이 논문은 대규모 비구조화 이산 공간에서의 베이지안 최적화 문제를 해결하기 위해, 사전 지식과 사후 확률 적응을 결합한 'ToSFiT(Thompson Sampling via Fine-Tuning)' 알고리즘을 제안하고, FAQ, 단백질 설계, 양자 회로 설계 등 다양한 작업에서 기존 방법론보다 우수한 샘플 및 계산 효율성을 입증합니다.

핵심

🌟 핵심 비유: "보물 찾기 게임과 천재 안내원"

상상해 보세요. 여러분은 거대한 미로 (아미노산 서열, 양자 회로, FAQ 답변 등) 안에 숨겨진 **최고의 보물 (가장 높은 점수를 주는 답)**을 찾아야 합니다. 문제는 미로가 너무 커서 모든 길을 다 걸어볼 수 없다는 것입니다. (예: 아미노산 조합의 수는 우주에 있는 원자 수보다 많습니다.)

기존의 방법들은 다음과 같은 문제가 있었습니다:

1. 지도 그리기 (기존 베이지안 최적화): 보물이 있을 법한 지역을 예측하는 '지도 (확률 분포)'를 그립니다. 그리고 그 지도에서 "가장 보물이 있을 확률이 높은 곳"을 찾기 위해 계산기를 두드려서 수학 문제를 풀어야 합니다.
   • 문제: 미로가 너무 크고 복잡하면, "어디가 최고일까?"를 계산하는 과정 자체가 너무 오래 걸려서 게임이 멈춰버립니다.

🚀 TOSFIT 의 해결책: "직관적인 안내원 (LLM) 을 훈련시키다"

이 논문은 수학 문제를 직접 풀지 않고, **천재 안내원 (대형 언어 모델, LLM)**을 고용하는 방식을 제안합니다.

1. 초기 상태 (초기화): 우리는 이미 엄청난 양의 책을 읽은 '초능력의 안내원 (사전 훈련된 LLM)'을 데려옵니다. 이 안내원은 미로의 구조를 대략적으로 알고 있습니다.
2. 훈련 (파인튜닝):
   • 안내원에게 "이쪽을 가보라"고 시키고, 보물을 찾았는지 확인합니다.
   • 보물을 찾았다면 "잘했다! 다음엔 이쪽으로 더 가보자"라고 칭찬하고, 찾지 못했다면 "아, 그쪽은 아니었구나"라고 가르칩니다.
   • 핵심: 이 과정에서 우리는 "어디가 최고일까?"를 수학적으로 계산하는 대신, 안내원의 머릿속 (모델의 가중치) 을 조금씩 수정하여, 안내원이 스스로 "가장 보물이 있을 확률이 높은 곳"을 찾아내도록 훈련시킵니다.

이것이 바로 TOSFIT입니다. 복잡한 계산을 생략하고, AI 가 스스로 학습하여 최적의 답을 생성하게 만드는 것입니다.

---

💡 이 기술이 왜 대단한가요? (3 가지 장점)

1. "계산기" 대신 "직관"을 사용합니다.

기존 방법은 "어디가 가장 좋을지"를 찾기 위해 모든 가능성을 계산해야 했기 때문에 컴퓨터가 과부하가 걸렸습니다. TOSFIT 은 AI 가 "아, 이쪽이 가장 유망해 보이네"라고 직관적으로 생성하게 합니다. 마치 수학 문제를 풀지 않고, 경험 많은 탐험가가 "저기 보물이 있을 것 같아"라고 바로 가리키는 것과 같습니다.

2. "실수"를 통해 빠르게 배웁니다.

이 기술은 **Thompson Sampling (톰슨 샘플링)**이라는 전략을 사용합니다.

• 비유: 주사위를 던져서 운에 맡기는 것이 아니라, "지금까지의 경험 (데이터) 을 바탕으로 보물이 있을 확률이 높은 곳"을 선택하는 것입니다.
• TOSFIT 은 이 확률을 AI 가 직접 학습하도록 하여, **적은 시도 (샘플)**로도 최고의 답을 찾아냅니다.

3. "실전"에서 proven 되었습니다.

연구진은 이 방법을 세 가지 어려운 분야에서 테스트했습니다.

• FAQ 답변 다듬기: 고객 문의에 대한 가장 적절한 답변을 찾아냅니다.
• 단백질 설계: 열에 강한 단백질을 찾아 의약품 개발에 도움을 줍니다. (우주의 원자 수보다 많은 조합 중 하나를 찾는 것!)
• 양자 회로 설계: 복잡한 양자 컴퓨터 회로를 설계합니다.

결과적으로, 기존 방법들 (진화 알고리즘, 강화학습 등) 보다 더 적은 시간과 계산 비용으로 더 좋은 결과를 냈습니다.

---

🎓 요약: 이 논문이 우리에게 주는 메시지

"거대한 미로에서 보물을 찾을 때, 모든 길을 계산하는 대신, 경험 많은 안내원 (AI) 을 고용하여 그 사람의 직관을 훈련시키는 것이 가장 빠르고 효율적이다."

이 연구는 AI 가 단순히 글을 쓰거나 그림을 그리는 것을 넘어, 과학적 발견과 복잡한 문제 해결의 핵심 도구로 자리 잡을 수 있음을 보여줍니다. 특히, 계산 비용이 너무 많이 들어 포기했던 문제들 (단백질 설계, 양자 컴퓨팅 등) 을 해결할 수 있는 열쇠를 쥐어주었습니다.

한 줄 요약:
"복잡한 수학 계산 대신, AI 를 '보물 찾기 전문가'로 훈련시켜, 적은 노력으로 최고의 답을 찾아내는 새로운 방법론입니다."

기술 요약

1. 문제 정의 (Problem)

대규모 비구조화 이산 공간 (Large Unstructured Discrete Spaces) 에서 베이지안 최적화 (Bayesian Optimization, BO) 를 수행할 때 발생하는 근본적인 한계가 있습니다.

• 기존 BO 의 한계: 전통적인 BO 는 불확실성 하에서 최적의 후보를 찾기 위해 '획득 함수 (Acquisition Function)'를 최대화하는 과정을 거칩니다. 연속 공간에서는 경사 하강법 (Gradient Ascent) 을 사용할 수 있지만, 아미노산 서열, 양자 회로 설계, 코드 생성과 같은 비구조화 이산 공간에서는 경사 (Gradient) 가 존재하지 않아 획득 함수를 효율적으로 최대화할 수 없습니다.
• 계산 비용: 이산 공간에서 획득 함수를 최대화하려면 모든 가능한 점을 탐색해야 하므로, 조합론적으로 매우 큰 공간 (예: 20 개의 아미노산으로 길이 100 의 서열) 에서는 계산이 불가능 (Intractable) 해집니다.
• 목표: 획득 함수의 명시적 최대화를 피하면서도, 베이지안 최적화의 강력한 수렴 보장 (Regret Bound) 을 유지하면서 대규모 이산 공간을 효율적으로 탐색할 수 있는 확장 가능한 방법론이 필요합니다.

2. 방법론 (Methodology: TOSFIT)

저자들은 **TOSFIT (Thompson Sampling via Fine-Tuning)**이라는 새로운 알고리즘을 제안합니다. 이는 생성형 대규모 언어 모델 (LLM) 의 미세 조정 (Fine-tuning) 을 통해 톰슨 샘플링 (Thompson Sampling) 을 구현하는 방식입니다.

• 핵심 아이디어:

  • 기존 톰슨 샘플링은 보상 사후분포에서 함수를 샘플링하고, 그 함수를 최대화하는 점을 선택합니다.
  • TOSFIT 는 획득 함수를 최대화하는 대신, **최대 보상을 얻을 확률 (Probability of Maximality, PoM)**을 직접 매개변수화 (Parameterize) 합니다.
  • 사전 훈련된 LLM 을 초기 정책 (Policy) 으로 사용하여, 프롬프트 조건 (Prompt-conditioning) 하에서 생성된 후보들이 PoM 분포를 따르도록 모델의 가중치를 미세 조정합니다.

• 알고리즘 흐름:

  1. 초기화: 사전 훈련된 LLM 을 기반으로 한 정책 $\pi_\theta$를 사용합니다.
  2. 후보 생성: 현재 정책으로 후보들을 생성하고 보상을 관측합니다.
  3. 보상 모델 업데이트: 관측된 데이터를 기반으로 가우시안 프로세스 (GP) 보상 모델을 업데이트합니다.
  4. VBOS 목적 함수 최적화: 변분 베이지안 낙관적 샘플링 (Variational Bayesian Optimistic Sampling, VBOS) 목적 함수를 사용하여 LLM 의 가중치를 미세 조정합니다.
     • VBOS 목적 함수는 보상 예측치 ($\mu_x$) 와 엔트로피 기반의 낙관적 탐색 보너스 ($\sqrt{-2 \ln \pi_x} \cdot \sigma_x$) 를 결합합니다.
     • 이를 통해 모델은 불확실성이 높은 영역을 탐색하면서도 높은 보상이 예상되는 영역을 활용 (Exploit) 합니다.
  5. 안정화: 고분산 (High Variance) 문제를 해결하기 위해 RLOO (Reinforce Leave-One-Out) 기법과 표준화된 어드밴티지 함수를 사용하여 학습을 안정화합니다.

• 확장성:

  • GP 추론을 선형 커널과 특징 맵 (Feature Map) 을 사용하여 폐쇄형 (Closed-form) 으로 수행함으로써, 관측 횟수에 의존하지 않는 상수 시간 복잡도 ($\Theta(d^2)$) 를 달성합니다.
  • 배치 (Batched) 최적화가 가능하여, 여러 후보를 병렬로 평가할 수 있습니다.

3. 주요 기여 (Key Contributions)

1. 새로운 Regret Bound (후회 상한) 유도:

   • 기존 VBOS 의 누적 후회 (Cumulative Regret) 상한을 $\tilde{O}(\sqrt{T|X|})$에서 $\tilde{O}(\sqrt{T\gamma_T})$로 개선했습니다. 여기서 $\gamma_T$는 최대 정보 획득량 (Maximum Information Gain) 으로, 커널 구조를 반영하여 조합론적으로 큰 공간에서도 유한한 상한을 가집니다.
   • 이 결과는 근사적인 (Gradient-based) VBOS 에 대해서도 일반화되었으며, 정책 초기화와 세심한 미세 조정의 중요성을 이론적으로 증명합니다.

2. TOSFIT 알고리즘 제안:

   • 사전 훈련된 LLM 의 지식을 활용하여 정책을 초기화하고, VBOS 목적 함수를 통해 사후 PoM 에 점진적으로 적응시키는 알고리즘을 설계했습니다.
   • 이는 획득 함수의 비효율적인 최대화를 우회하면서도 톰슨 샘플링의 강력한 수렴 보장을 유지합니다.

3. 다양한 도메인에서의 SOTA 성능 입증:

   • FAQ 응답 개선 (Natural Language): Qwen3 모델을 사용하여 질문 답변의 품질을 최적화.
   • 단백질 탐색 (Thermally Stable Protein): ProtGPT2 를 사용하여 열적 안정성이 높은 아미노산 서열 설계. (검색 공간이 관측 가능한 우주의 원자 수를 초과함).
   • 양자 회로 설계 (Quantum Circuit Design): Qwen2.5-Coder 를 사용하여 저에너지 양자 상태를 준비하는 회로 생성.
   • 위 세 가지 작업에서 기존 베이지안 최적화, 강화학습 (Actor-Critic), 진화적 탐색 (Evolutionary Search) 방법론 대비 **샘플 효율성 (Sample Efficiency)**과 계산 효율성 (Computational Efficiency) 모두에서 최첨단 (State-of-the-Art) 성능을 보였습니다.

4. 실험 결과 (Results)

• 샘플 효율성: TOSFIT 는 제한된 평가 횟수 내에서 가장 높은 보상을 발견하는 속도가 가장 빠릅니다. 특히 'Unguided Generation'이나 'Actor-Critic' 기반 방법들은 초기에 성능이 포화되거나 탐색/활용의 균형을 못 맞추는 반면, TOSFIT 는 낙관적 탐색 (Optimism) 을 통해 지속적으로 개선됩니다.
• 계산 효율성: 획득 함수 최대화를 위한 반복적 탐색이 필요 없어, 전체적인 계산 비용 대비 성능이 뛰어납니다. 특히 배치 처리 (Batching) 시에도 효율적입니다.
• 모델 크기와 사전 지식의 중요성:
  • 더 큰 모델 (8B 파라미터) 을 사용할수록 성능이 향상되지만, TOSFIT 는 작은 모델 (0.6B~1.7B) 에서도 강력한 사전 지식 (Pre-training) 을 통해 우수한 성능을 발휘합니다.
  • 주의: 학습률 (Learning Rate) 이 너무 크면 사전 지식을 잊어버리고 (Catastrophic Forgetting) 성능이 저하되며, 너무 작으면 적응이 느려집니다. 따라서 **신중한 미세 조정 (Careful Fine-tuning)**이 필수적입니다.
• 다양성 유지: TOSFIT 는 엔트로피 정규화를 통해 정책의 다양성 붕괴 (Diversity Collapse) 를 방지하고, 불확실성이 높은 영역에서는 탐색을, 확신이 높은 영역에서는 활용을 자연스럽게 조절합니다.

5. 의의 및 결론 (Significance)

• 이산 공간 최적화의 패러다임 전환: 획득 함수 최대화라는 계산적 병목 현상을 제거하고, 생성형 AI 의 생성 능력을 직접 최적화 도구로 활용하는 새로운 접근법을 제시했습니다.
• 이론과 실증의 결합: 톰슨 샘플링의 강력한 이론적 보장 (Regret Bound) 을 LLM 미세 조정이라는 실용적인 기법과 성공적으로 결합했습니다.
• 과학적 발견 가속화: 단백질 설계, 양자 컴퓨팅, 화학 물질 발견 등 기존에 탐색이 불가능했던 거대 이산 공간에서 AI 를 활용한 자동화된 과학적 발견을 가능하게 하는 강력한 도구가 될 것으로 기대됩니다.

요약하자면, TOSFIT는 LLM 의 생성 능력을 베이지안 최적화의 프레임워크에 통합하여, 획득 함수 최대화의 비효율성을 해결하고 대규모 이산 공간에서 효율적이고 강력한 최적화를 달성한 획기적인 연구입니다.