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

The paper introduces COREA, a cost-efficient reasoning system that cascades a small language model with a large language model and employs reinforcement learning to calibrate confidence scores, thereby significantly reducing inference costs while maintaining high accuracy.

The Gist

Imagine you run a busy, high-end restaurant. You have two chefs:

1. The "Speedy Sous-Chef" (SLM): This chef is fast, cheap to employ, and great at cooking simple dishes like grilled cheese or salads. However, they sometimes get confused by complex recipes (like a 10-course molecular gastronomy meal) and might confidently serve you a burnt dish.
2. The "Master Chef" (LLM): This chef is a genius. They can cook anything perfectly. But they are incredibly expensive, slow, and their time is limited. You can't afford to have them cook every single order.

The Problem:
In the world of AI, we want to solve hard problems (like math or science) using the Master Chef, but it costs too much money to use them for everything. If we use the Speedy Chef for everything, we save money, but we get wrong answers.

The Solution: COREA (The Smart Waiter System)
The authors of this paper created a system called COREA. Think of it as a super-smart waiter who stands between the customer and the kitchen.

Here is how it works, step-by-step:

1. The "Confidence Check"

When a customer orders a dish (a question), the Speedy Sous-Chef tries to cook it first. But here is the magic trick: The Speedy Chef is trained to know what they don't know.

Before serving the dish, the Speedy Chef has to say out loud: "I am 90% sure this is perfect" or "I am only 20% sure, this looks risky."

• Old Way: The Speedy Chef would often say, "I'm 100% sure!" even when they were wrong. This is called being "overconfident."
• The COREA Way: Through a special training process (Reinforcement Learning), the Speedy Chef learns to be honest. If the dish is hard, they admit, "I'm not confident."

2. The Decision Gate

The Smart Waiter (the system) listens to that confidence score:

• If the Chef says "I'm confident (above 70%)": The Waiter serves the dish immediately. Result: You get a fast, cheap answer.
• If the Chef says "I'm not confident (below 70%)": The Waiter says, "Hold on, this is too tricky." They hand the order over to the Master Chef. Result: You get a perfect answer, but it costs more.

3. The Training (The "Taste Test")

How did they teach the Speedy Chef to be honest?
They didn't just tell them to be nice. They used a "Reward System" during training:

• Reward for being right: If the answer is correct, the chef gets a point.
• Reward for being honest: If the chef says "I'm 50% sure" and they are actually right 50% of the time, they get a bonus. If they say "I'm 100% sure" but get it wrong, they get a penalty.

This forced the Speedy Chef to align their confidence with their actual ability. They learned to say "I don't know" when they truly didn't know.

The Results: A Win-Win

The paper tested this system on thousands of math and logic problems. Here is what happened:

• Cost Savings: By letting the Speedy Chef handle the easy stuff (which is most of the time), they saved 16% to 21% of the money compared to using the Master Chef for everything.
• Accuracy: The system was still almost as accurate as using the Master Chef alone (only about 2% less accurate).
• Efficiency: It's like having a team where the junior staff handles 60% of the work, and the senior staff only steps in for the hard 40%.

The Analogy Summary

Imagine you are a student taking a test.

• Without COREA: You either use a calculator for every single math problem (expensive/slow) or you guess on everything (fast but wrong).
• With COREA: You try to solve the problem in your head first. If you feel confident, you write down the answer. If you feel stuck or unsure, you immediately raise your hand and ask the teacher (the Master Chef) for help.

This paper proves that if you train your "student" (the small AI) to be honest about their own knowledge, you can build a system that is cheap, fast, and smart all at once.

Technical Summary

Here is a detailed technical summary of the paper "Confidence-Calibrated Small-Large Language Model Collaboration for Cost-Efficient Reasoning" (COREA).

1. Problem Statement

Large Language Models (LLMs) demonstrate superior reasoning capabilities (e.g., in mathematics and coding) compared to Small Language Models (SLMs), primarily due to their ability to generate explicit Chain-of-Thought (CoT) reasoning traces. However, this capability comes with prohibitively high inference costs and latency.

Existing solutions face a trade-off:

• SLMs alone: Cost-effective but lack reasoning robustness and often exhibit overconfidence (they do not know what they don't know), leading to poor accuracy on complex tasks.
• LLMs alone: High accuracy but high operational cost.
• Current Collaboration/Routing: Methods that route queries between models often rely on external classifiers, heuristics, or token-probability averages, which fail to accurately capture the model's internal reasoning confidence or require additional sampling overhead.

Core Challenge: How to enable an SLM to accurately recognize its own limitations (self-awareness) and dynamically defer difficult problems to an LLM, thereby balancing accuracy and cost without external routing modules.

2. Methodology: COREA Framework

The authors propose COREA (COllaborative REAsoner), a cascaded system where an SLM acts as the first line of defense, deferring only uncertain queries to a more powerful LLM.

A. System Architecture

1. Inference Flow: For a given query, the SLM is prompted to generate:
   • Reasoning steps.
   • A final answer.
   • A verbalized confidence score ($y_c \in [0, 1]$).
2. Decision Logic:
   • If $y_c \geq T$ (a predefined threshold, typically the LLM's baseline accuracy), the SLM's answer is accepted.
   • If $y_c < T$, the query is deferred to the LLM for a final answer.

B. Training: Reinforcement Learning with Confidence Calibration (RLCC)

The core innovation is training the SLM not just for correctness, but for calibrated confidence. The authors use Group Relative Policy Optimization (GRPO) with a composite reward function:

$$R = R_{correct} + R_{format} + R_{confidence}$$

1. Correctness Reward ($R_{correct}$): Binary reward ($1$if answer matches ground truth,$0$ otherwise).
2. Format Reward ($R_{format}$): Ensures the model outputs reasoning, the answer in a specific box, and the confidence score in the required format.
3. Confidence Reward ($R_{confidence}$): This is the novel component. It penalizes the distance between the model's predicted confidence ($y_c$) and the true probability of correctness ($p$).
   • Since $p$ is unknown during inference, it is estimated during training using group accuracy ($\hat{p}$) from $N$ sampled responses for the same question.
   • The paper evaluates several distance metrics (L1, L2, KL) and finds that the L1 distance ($R_{L1} = -|p - y_c|$) offers the best balance.

Key Distinction: Unlike previous methods that estimate confidence based on a single sample's correctness, COREA estimates the "true" probability $p$ based on the group's performance, encouraging the model to output a confidence score that reflects the expected accuracy of that specific question type.

3. Key Contributions

1. COREA Framework: A practical SLM-LLM collaboration system that achieves high accuracy with significantly reduced costs by leveraging the SLM's self-calibrated confidence to decide when to defer.
2. RLCC Training Algorithm: A novel reinforcement learning approach that simultaneously improves reasoning ability and calibrates confidence scores. It introduces a group-level confidence reward that aligns the SLM's verbalized confidence with its actual correctness probability.
3. Comprehensive Evaluation: Extensive experiments across diverse datasets (Math, Science, Commonsense) and model backbones (Qwen, Llama) demonstrating that the method generalizes well.

4. Experimental Results

The authors evaluated COREA using Qwen2.5-7B (SLM) and Qwen2.5-32B (LLM) on multiple benchmarks (DeepMath, Math500, GSM8K, GPQA, CommonsenseQA).

• Cost Efficiency: Compared to using the LLM alone, COREA reduced inference costs by:
  • 21.5% on Out-of-Domain (OOD) Math datasets.
  • 16.8% on OOD Non-Math datasets.
  • 6.7% on In-Domain Math (DeepMath500).
• Accuracy Preservation: The system maintained a Pass@1 score within 2% of the standalone LLM baseline, a significant improvement over other collaborative methods which often suffered larger accuracy drops (e.g., 5–10% drops) when trying to save costs.
• Calibration Performance:
  • The proposed L1-SLM achieved the lowest Expected Calibration Error (ECE) (0.12) and highest AUROC (0.72) among all tested methods.
  • Standard RL training (RLVR) improved reasoning but failed to calibrate confidence (models remained overconfident).
• Ablation Studies:
  • Reward Functions: L1 reward outperformed L2 and KL in balancing accuracy and calibration.
  • Model Sizes: The method was effective across different SLM sizes (1.5B, 7B, 8B), though smaller models (1.5B) showed slightly less self-awareness than larger ones.
  • Rollout Size: The method is robust to rollout sizes; even a small rollout size (4 or 8) yielded strong results, making training efficient.

5. Significance and Impact

• Practical Deployment: COREA provides a viable path for deploying high-reasoning AI systems in cost-sensitive real-world scenarios. It eliminates the need for expensive external routers or complex architectural changes.
• Solving the "Overconfidence" Problem: The paper addresses a critical flaw in current LLMs/SLMs—their inability to know when they are wrong. By explicitly training for confidence calibration, the system creates a "self-aware" agent that knows when to ask for help.
• Scalability: The approach is model-agnostic and can be applied to various SLM-LLM pairings, offering a generalizable strategy for the "Reasoning Economy."

In conclusion, COREA demonstrates that calibrated confidence is a powerful mechanism for enabling efficient collaboration between models of different scales, achieving near-optimal accuracy while drastically reducing computational costs.