Model Cascade: making LLM classification cheaper
Many LLM workloads are classification tasks. This can get expensive, and it is going to become more and more important, especially with the proliferation of software factories. So what is Model Cascade? In short, it is a way to build a deterministic system around a cheap model and make it give us the same results as the expensive model.
Core concepts
- A Proxy is the cheap model. It returns an output and a confidence score.
- An Oracle is the expensive model. It returns its own output and whether the proxy output was correct.
- BARGAIN_A is the accuracy target mode: match the oracle on at least the target percent of records, using the proxy as often as possible.
- BARGAIN_P and BARGAIN_R are precision and recall target modes for binary tasks, with a fixed oracle call budget.
The principle
The LLM we use gives us the probability of every token in the output, using the same probability model used to generate the response. We put all the tokens of the response together, and we get the probability of the response.
Now the smart part of the Model Cascade:
- We do a sample run with the Oracle model for, letβs say, 500-1000 samples.
- After that, we do the same with the Proxy model, and we get the probability numbers.
- We filter the ones that matched the classification of the oracle, so we get a range in which we know our Proxy model should be correct.
- We can check that range when we are doing the classification.
flowchart TB
subgraph CAL["Calibrate once, offline"]
S["Sample ~500 records"] --> O1["Label sample with oracle"]
O1 --> T["Try every observed confidence <br/> value as a threshold"]
T --> P["Pick cheapest threshold that<br/>meets the accuracy target"]
O1 --> G["Check that proxy confidence agrees <br/> with oracle labels"]
end
subgraph ROUTE["Route every record, at scale"]
R["Record"] --> PX["Proxy: small, cheap model"]
PX --> L["Label + confidence score,<br/>from logprob"]
L --> D{"Confidence above threshold?"}
D -->|"yes, most records"| K["Keep proxy label"]
D -->|"no, few records"| O2["Oracle: large, expensive model"]
K --> OUT["Final labels"]
O2 --> OUT
end
P -. "sets threshold" .-> D
The BARGAIN Paper
Below is a summary of the BARGAIN paper I used to learn about this principle. It is more detailed than the first part, so if you want to learn more, read on. Or read the full paper here: https://github.com/ucbepic/BARGAIN
What BARGAIN reports
Across eight datasets, the BARGAIN paper reports up to 86% more cost reduction than competing methods.
The follow-up Task Cascades paper adds three optimizations: rewriting prompts into simpler surrogate questions, reading only the most relevant document chunks, and searching over candidate cascades for the cheapest sequence. These cut costs a further 48.5% on average.
Unlike FrugalGPT, BARGAIN gives statistical guarantees. Unlike SUPG, the guarantees hold at any sample size, and BARGAIN uses adaptive sampling and better estimation.
Using the BARGAIN library
pip install bargain
Dependencies are numpy, pandas, tqdm, and openai. You can swap providers by defining your own proxy and oracle.
Reference points from the repo examples
Examples live in examples/. Run the Supreme Court one from that directory; it loads court_opinion.csv by relative path.
- Toy binary task: accuracy 0.95, proxy used on 45% of records
- Open-ended extraction: accuracy 1.0, proxy used on 57% of records
- Supreme Court opinions: accuracy 0.976, proxy used on 40.6% of records (
target=0.9,delta=0.1)
The Supreme Court numbers come from one run and may change with model versions, API behavior, or dataset changes.
A practical order of work
- Pick your oracle and proxy, and confirm the proxy can provide a useful confidence score.
- Write the task prompt. Use
True/Falseonly for binary tasks. - Run
BARGAIN_Aon a sample with your target and delta to see what fraction the proxy can handle. - If that fraction is low, the logits do not track the oracle. Try a different proxy or a simpler prompt before tuning anything else.
- For extra savings, apply the Task Cascades ideas: surrogate questions and relevant chunks only.
Getting logprobs with LangChain OpenAI models
Pass logprobs and top_logprobs to ChatOpenAI, then read the scores from response_metadata:
import math
from langchain_openai import ChatOpenAI
llm = ChatOpenAI(
model="gpt-5-nano",
temperature=0,
logprobs=True,
top_logprobs=5,
)
response = llm.invoke(
"Does the text 'zebra' mention an animal? Answer with only True or False."
)
content = response.response_metadata["logprobs"]["content"]
first_token = content[0]
print(first_token["token"], first_token["logprob"]) # e.g. "True" -0.01
print(math.exp(first_token["logprob"])) # probability, e.g. 0.99
Each entry in content is one token with its own logprob. The snippet reads only the first token, which works because the prompt forces a single-word answer. For a multi-token answer, sum all token logprobs instead:
total_logprob = sum(t["logprob"] for t in content)
For classification, prompt for a single word so the response is one token, then use that tokenβs logprob as the confidence score.
Getting the score for a specific label
The top token is the
Comments
No comments yet. Start the discussion.