Token-Budget-Aware LLM Reasoning

Reasoning is critical for large language models (LLMs) to excel in a wide range of tasks. While methods like Chain-of-Thought (CoT) reasoning and enhance LLM performance by decomposing problems into intermediate steps, they also incur significant overhead in token usage, leading to increased costs. We find that the reasoning process of current LLMs is unnecessarily lengthy and it can be compressed by including a reasonable token budget in the prompt, but the choice of token budget plays a crucial role in the actual compression effectiveness. We then propose a token-budget-aware LLM reasoning framework that dynamically adjusts the number of reasoning tokens based on the reasoning complexity of each problem. Experiments show that our method effectively reduces token costs in CoT reasoning with only a slight performance reduction, offering a practical solution to balance efficiency and accuracy in LLM reasoning. Code: https://github.com/GeniusHTX/TALE

한 줄 요약

TALE은 CoT 추론 과정에서 토큰 비용을 67%까지 줄이며 정확도를 3% 미만으로 유지하는 토큰 예산 인식 추론 프레임워크이다.

핵심 기여도

TALE-EP는 Zero-shot 프롬프팅을 통해 문제별 토큰 예산을 추정하고 CoT 추론에 통합하여 67%의 토큰 사용 감소를 달성.
TALE-PT는 토큰 예산 인식을 사후 학습을 통해 내재화하여 CoT 대비 약 50%의 토큰 사용 감소.
토큰 예산 설정에 따른 "Token Elasticity" 현상을 분석하고, 최적 예산을 탐색하는 방법 제시.
GSM8K, GSM8K-Zero, MathBench 데이터셋에서 실험적으로 검증.

핵심 아이디어

기존 CoT 추론은 문제를 단계별로 분해하여 정확도를 향상시키지만, 이 과정에서 많은 토큰이 소비되어 비용 문제가 발생한다. 연구자들은 CoT 추론 과정이 불필요하게 길고, 토큰 예산을 프롬프트에 포함하면 추론을 압축할 수 있음을 발견했다. 그러나 토큰 예산의 선택이 압축 효과에 큰 영향을 미친다는 점에서 "Token Elasticity" 현상이 관찰된다. 즉, 예산이 작을 경우 LLM이 예산을 따르지 못하고, 실제 토큰 사용량이 예산보다 훨씬 증가한다. 이를 해결하기 위해 TALE은 문제 복잡도에 따라 토큰 예산을 동적으로 조정하는 프레임워크를 제안한다. TALE-EP는 Zero-shot 프롬프팅을 기반으로 예산을 추정하고, TALE-PT는 사후 학습을 통해 토큰 예산 인식을 내재화한다.

기술적 접근법

**모듈명**: TALE-EP (Token Budget Estimation and Prompting), TALE-PT (Token Budget Awareness Internalization via Post-training)
**프롬프팅 기반 추정**: TALE-EP는 Zero-shot 프롬프팅을 사용해 문제별 토큰 예산을 추정.
**사후 학습**: TALE-PT는 토큰 예산 인식을 모델 내부에 학습시켜, 명시적 예산 없이도 효율적인 추론 가능.
**데이터셋**: GSM8K, GSM8K-Zero, MathBench
**모델**: GPT-4o, GPT-4o-mini, Yi-lightning, o3-mini, Llama-3.1-8B-Instruct
**메트릭**: 정확도 (Accuracy), 출력 토큰 수 (Number of Output Tokens)

주요 결과

**GSM8K-Zero 데이터셋**에서 TALE-EP는 CoT 대비 67%의 토큰 사용 감소를 달성하면서 정확도는 3% 미만 감소.
**TALE-PT**는 CoT 대비 약 50%의 토큰 사용 감소를 보이며 경쟁력 있는 성능 유지.
**Token Elasticity** 현상: 50 토큰 예산 시 258 토큰 → 86 토큰 감소, 10 토큰 예산 시 157 토큰 → 효과적 압축 실패.

의의 및 한계

TALE은 CoT 추론의 토큰 비용 문제를 해결하며, 효율성과 정확도를 균형 있게 유지하는 실용적 해결책을 제시한다. 특히, Zero-shot 프롬프팅과 사후 학습을 결합한 이중 접근법은 다양한 LLM에 적용 가능하다는 장점이 있다. 그러나 토큰 예산 설정이 문제 복잡도에 따라 달라지므로, 모든 문제에 대해 동일한 예산을 적용할 수 없다는 한계가 있다. 또한, 사후 학습 기반의 TALE-PT는 추가적인 학습 비용이 발생할 수 있다.

실용적 활용

TALE은 클라우드 기반 LLM 서비스, 대규모 추론 작업, 비용 민감한 응용 분야에서 유용하게 활용될 수 있다. 특히, 토큰 비용을 줄이면서도 정확도를 유지해야 하는 산업 현장에서 효과적일 것으로 기대된다.