Table of Contents
Fetching ...

Confidence-Calibrated Small-Large Language Model Collaboration for Cost-Efficient Reasoning

Chuang Zhang, Zizhen Zhu, Yihao Wei, Bing Tian, Junyi Liu, Henan Wang, Xavier Wang, Yaxiao Liu

TL;DR

This work proposes COllaborative REAsoner, a system that cascades an SLM with an LLM to achieve a balance between accuracy and cost in complex reasoning tasks, and introduces a reinforcement learning-based training algorithm that aligns the SLM's confidence through an additional confidence calibration reward.

Abstract

Large language models (LLMs) demonstrate superior reasoning capabilities compared to small language models (SLMs), but incur substantially higher costs. We propose COllaborative REAsoner (COREA), a system that cascades an SLM with an LLM to achieve a balance between accuracy and cost in complex reasoning tasks. COREA first attempts to answer questions using the SLM, which outputs both an answer and a verbalized confidence score. Questions with confidence below a predefined threshold are deferred to the LLM for more accurate resolution. We introduce a reinforcement learning-based training algorithm that aligns the SLM's confidence through an additional confidence calibration reward. Extensive experiments demonstrate that our method jointly improves the SLM's reasoning ability and confidence calibration across diverse datasets and model backbones. Compared to using the LLM alone, COREA reduces cost by 21.5% and 16.8% on out-of-domain math and non-math datasets, respectively, with only an absolute pass@1 drop within 2%.

Confidence-Calibrated Small-Large Language Model Collaboration for Cost-Efficient Reasoning

TL;DR

This work proposes COllaborative REAsoner, a system that cascades an SLM with an LLM to achieve a balance between accuracy and cost in complex reasoning tasks, and introduces a reinforcement learning-based training algorithm that aligns the SLM's confidence through an additional confidence calibration reward.

Abstract

Large language models (LLMs) demonstrate superior reasoning capabilities compared to small language models (SLMs), but incur substantially higher costs. We propose COllaborative REAsoner (COREA), a system that cascades an SLM with an LLM to achieve a balance between accuracy and cost in complex reasoning tasks. COREA first attempts to answer questions using the SLM, which outputs both an answer and a verbalized confidence score. Questions with confidence below a predefined threshold are deferred to the LLM for more accurate resolution. We introduce a reinforcement learning-based training algorithm that aligns the SLM's confidence through an additional confidence calibration reward. Extensive experiments demonstrate that our method jointly improves the SLM's reasoning ability and confidence calibration across diverse datasets and model backbones. Compared to using the LLM alone, COREA reduces cost by 21.5% and 16.8% on out-of-domain math and non-math datasets, respectively, with only an absolute pass@1 drop within 2%.
Paper Structure (30 sections, 41 equations, 9 figures, 15 tables)

This paper contains 30 sections, 41 equations, 9 figures, 15 tables.

Figures (9)

  • Figure 1: (a) Architecture of COREA:COREA cascades an SLM with an LLM. For each query, the SLM is prompted to generate reasoning steps, an answer, and a confidence score. LLM is invoked only when the SLM's confidence is less than the predefined threshold. (b) Confidence calibration of the SLM: The SLM is trained using GRPO with a multi-objective reward function to improve its reasoning ability and calibrate its confidence.
  • Figure 2: Results on DeepMath-500: (a) Pass@1, (b) average cost, and (c) LLM usage percentage versus confidence threshold; (d) Pass@1 versus average cost for different systems. The vertical dashed lines in (a)-(c) mark the Baseline LLM's Pass@1 (0.69).
  • Figure 3: Training process of Qwen2.5-7BI with different reward configurations: RLVR, Brier (RLCC with Brier score), and L1 (RLCC with L1 confidence reward) on DeepMath500, with 8 sampled answers per question.
  • Figure 4: Comparison of Pass@1 and ECE across several models trained with different reward configurations, averaged across all datasets. Detailed results on each dataset are provided in Appendix \ref{['sec:detailed_calibration']}.
  • Figure 5: Question and Response of Qwen2.5-7BI.
  • ...and 4 more figures