CoLT: Reasoning with Chain of Latent Tool Calls

기술 요약: CoLT – 잠재적 도구 호출 체인을 통한 추론 (Reasoning with Chain of Latent Tool Calls)

문제 정의

Chain-of-Thought (CoT) 프롬프팅은 복잡한 문제를 순차적인 단계로 나누어 대규모 언어 모델(LLM)의 추론 능력을 향상시키는 데 필수적임이 입증되었습니다. 그러나 명시적인 CoT는 추론 체인을 토큰 단위로 생성해야 하므로 추론 과정에서 높은 계산 비용을 초도합니다. 이를 해결하기 위해 연속적인 잠재 공간(latent space) 내에서 추론하여 토큰 길이를 줄이는 암묵적 CoT 방법들(예: COCONUT, CODI, COLAR)이 제안되었으나, 이들은 일반적으로 두 가지 결정적인 한계를 가집니다:

1. 구조적 경직성: 이들은 종종 상당한 모델 아키텍처 수정과 철저한 학습을 요구합니다.
2. 적용 가능성: 이러한 제약 사항은 다양한 작업 및 모델 구조에 대한 광범위한 채택을 제한합니다.

또한, 기존의 잠재 추론 방식은 완전히 "블랙박스" 형태의 잠재 공간 내에서 작동하는 경우가 많아, 베이스 모델의 해석 가능성과 사전 학습된 추론 능력을 희생할 잠재적 위험이 있습니다.

방법론: CoLT 프레워크

저자들은 잠재 추론을 "파라미터화된 도구 호출(parametric tool calls)"로 구현하는 새로운 프레임워크인 **CoLT (Chain-of-Latent-Tools)**를 제안합니다. CoLT는 메인 LLM이 완전히 잠재 공간 내에서 추론하는 대신, 명시적인 토큰 공간에서 작동하면서 추론 단계의 확장을 외부의 미분 가능한 디코더(differentiable decoders)에 위임하도록 합니다.

핵심 메커니��ㅁ

1. 시드 토큰(Seed Tokens) 및 잠재적 도구 호출:
   추론 과정 중에 메인 LLM은 압축된 추론 단계를 나타내는 특수 "시드 토큰"을 생성합니다. 이는 다음으로 구성됩니다:

   • 바디 토큰 (<BDY>): 추론 단계의 압축된 정보를 담고 있는 잠재 임베딩의 운반체입니다.
   • 트리거 토큰 (<TRG>): 잠재적 도구 호출의 종료를 알리고 어떤 외부 디코더를 호출할지 지정하는 지표입니다.

   전형적인 잠재 도구 호출 시퀀스는 다음과 같습니다: [질문] [이전 단계들] <BDY> ... <TRG>.

2. 미분 가능한 디코딩 (Differentiable Decoding):
   트리거 토큰이 발생하면, 시드 토큰의 은닉 상태(hidden states)가 추출되어 외부 미분 가능한 디코더(예: 경량 트랜스포머, RNN 또는 Multi-hot 디코더)로 전달됩니다.

   • 디코더는 은숨 상태를 다시 완전하고 명시적인 텍스트 토큰(예: 수학 방정식)으로 풀어냅니다(unpack).
   • 이렇게 디코딩된 토큰들은 컨텍스트에 결합되며, 메인 LLM은 업데이트된 읽기 가능한 컨텍스트와 함께 자기회귀적(autoregressive) 생성을 계속합니다.

3. 학습 목적 함수 (Training Objectives):
   이 프레임워크는 두 가지 지도 학습 손실 구성 요소를 사용하여 엔드투엔드(end-to-end)로 학습됩니다:

   • 메인 손실 ($L_{main}$): 메인 모델이 시드 토큰과 트리거 신호를 올바르게 생성하도록 보장합니다.
   • 잠재 손실 ($L_{lat}$): 디코더가 시드 임베딩으로부터 텍스트 토큰을 올바르게 재구성하도록 보장합니다.
   • 총 손실: $L_{sup} = L_{main} + L_{lat}$.

4. 강화 학습(RL) 호환성:
   디코더가 미분 가능하고 샘플링을 지원할 수 있기 때문에, 전체 추론 과정은 다회차 대화(multi-round conversation)로 모델링될 수 있습니다. 이를 통해 그룹 상대 정책 최적화(GRPO)를 적용하여, 수동으로 노이즈를 주입하지 않고도 다양한 추론 경로를 탐색하고 올바른 체인을 강화할 수 있습니다.

주요 기여

1. 새로운 프레임워크: CoLT는 잠재 추론을 도구 호출로 구현하는 방법을 도입하여, 메인 모델이 시드 토큰으로 중간 단계를 압축하면서도 사전 학습된 텍스트 공간 내에서 추론할 수 있게 합니다.
2. 디코더 유연성: 저자들은 다양한 디코더 구조(Transformer, RNN, Multi-hot)를 설계하고 평가하였으며, 트랜스포머가 가장 우수한 성능을 보이지만 프레임워크가 다양한 형태에 적응 가능하다는 것을 입증했습니다.
3. 효율성 및 성능: CoLT는 기존 잠재 모델들에 비해 더 높은 정확도와 더 짧은 추론 길이를 달者히합니다. 또한, 더 적은 학습 에포크와 적은 구조적 수정을 요구합니다.
4. RL 통합: 이 프레임워크는 강화 학습 알고리즘과 호환되는 것으로 나타났으며, 탐색을 통해 어려운 데이터셋에서 성능 향상을 이끌어낼 수 있음을 보여주었습니다.

실험 결과

저자들은 네 가지 수학적 추론 벤치마크인 GSM8k, GSM8k-hard, SVAMP, MultiArith에서 CoLT를 평가했습니다.

• 정확도 vs. 길이: CoLT는 정확도와 추론 체인 길이 모두에서 기존 잠재 추론 베이스라인(COCONUT, CODI, COLAR, SIM-CoT 포함)을 능가했습니다. 예를 들어, GSM8k-Aug에서 CoLT는 COLAR(2x) 대비 5%의 정확도 향상을 달성하는 동시에, 추론 토큰 길이를 12.7에서 7.73으로 단축했습니다.
• 일반화: 모델은 특히 MultiArith 데이터셋에서 지도 학습 기반의 CoT 베이스라인을 능가하며 강력한 도메인 외 일반화 능력을 보여주었습니다.
• 절제 연구 (Ablation Studies):
  • 디코더 레이어: 디코더 레이어와 시드 토큰 길이를 늘리면 정확도가 향상되지만, 계산 비용 대비 수익 체감(diminishing returns)이 발생합니다.
  • 디코더 구조: 트랜스포머 디코더가 가장 좋은 성능을 보였고, 그 뒤를 RNN이 이었습니다. Multi-hot 디코더는 성능이 저조했는데, 이는 패킹된 시드를 풀어내는 데에는 순환 구조(recurrent structures)가 더 적합함을 시사합니다.
• 강화 학습: GSM8k-Aug 및 MATH 데이터셋(DeepSeek-R1-Distill-Qwen-1.5B 사용)에 대해 CoLT 모델에 GRPO를 적용했을 때 성능이 더욱 향상되었으며, 이는 다양한 추론 경로로부터 학습할 수 있는 프레임워크의 능력을 증명합니다.

의의 및 주장

본 논문은 CoLT가 추론 효율성과 모델 능력 사이의 트레이드오프에 대한 실질적인 해결책을 제공한다고 주장합니다. 메인 모델을 사전 학습된 텍스트 공간에 고정하고 외부 디코더를 사용하여 "잠재적" 확장을 처리함으로써, CoLT는 추론 체인의 해석 가능성(최종 출력은 완전한 텍스트 형태임)을 유지하면서도 추론 비용을 크게 줄입니다.

저자들은 CoLT가 원래 모델 구조를 최소한으로 수정하면서도 적용 가능하다는 점을 강조하며, 이는 과도한 아키텍처 변경을 요구하는 이전의 잠재 추론 방식들보다 더 폭넓게 적용될 수 있음을 의미합니다. 또한, CoLT의 강화 학습 호환성은 도구 호출의 세밀함(granularity)과 디코더 구조를 특정 도메인에 맞게 적절히 조정한다면, 잠재 추론을 더 복잡한 작업으로 확장할 수 있는 경로를 제시합니다. 논문은 현재 수학적 추론에 집중되어 있지만, 잠재적 도구 호출의 개념이 엔티티 검색(entity retrieval) 및 멀티모달 추론을 포함한 더 넓은 응용 분야를 가질 잠재력이 있다고 결론짓습니다.