Model Cascade: making LLM classification cheaper

Many LLM workloads are classification tasks. This can get expensive, and I believe 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 make 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, 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 lets say 500-1000 samples
• After that we do same with Proxy model (cheap one), we get the probability numbers
• We filter the ones that matched the classification of the oracle so we get a range we know that 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, they hold at any sample size, and it 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

1. Pick your oracle and proxy, and confirm the proxy can provide a useful confidence score.
2. Write the task prompt. Use True/False only for binary tasks.
3. Run BARGAIN_A on a sample with your target and delta to see what fraction the proxy can handle.
4. If that fraction is low, the logits do not track the oracle. Try a different proxy or a simpler prompt before tuning anything else.
5. 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 model's answer. For a label it did not pick, look inside top_logprobs:

candidates = {c["token"]: c["logprob"] for c in first_token["top_logprobs"]}
score_for_true = candidates.get("True")

If a label is absent from top_logprobs, its score is unavailable. Do not treat a fallback value as the model's actual score.

Wiring it into BARGAIN's proxy_func

def proxy_func(self, data_record: str):
    response = llm.invoke(self.task.format(data_record))
    first = response.response_metadata["logprobs"]["content"][0]
    return first["token"], first["logprob"]

For binary classification, request enough top_logprobs entries to include both labels, normalize the two label probabilities, and return the probability of the selected label:

import math

def proxy_func(self, data_record: str):
    response = llm.invoke(self.task.format(data_record))
    first = response.response_metadata["logprobs"]["content"][0]
    candidates = {
        item["token"]: math.exp(item["logprob"])
        for item in first["top_logprobs"]
    }
    true_prob = candidates.get("True", 0.0)
    false_prob = candidates.get("False", 0.0)
    total = true_prob + false_prob
    if not total:
        return False, 0.0
    true_prob /= total
    false_prob /= total
    output = true_prob > false_prob
    return output, true_prob if output else false_prob

Gotchas

• Set temperature=0 for more repeatable answers. It does not guarantee identical responses, logprobs are not calibrated probabilities of correctness, and some reasoning models disallow temperature.
• Not every provider exposes logprobs. Anthropic does not, so a Claude proxy needs another confidence signal, such as a judge call. OpenAI and most open-weight models served with vLLM or HuggingFace do support them.
• If response_metadata has no logprobs, the provider did not return them. logprobs and top_logprobs are direct ChatOpenAI arguments; other provider-specific parameters go in extra_body.