Cut Classification Costs With a Model Cascade
This note covers Shreya Shankar’s session in the AI Product Engineering series.
Many LLM workloads are classification tasks, a problem that predates generative AI. LLMs are so easy to use that it is tempting to throw your largest model at everything. However, when you classify at scale this becomes expensive.
Shreya’s final talk teaches BARGAIN, a method for shifting classification work to a smaller model while meeting an accuracy target you choose. The method uses the probability the small model emits with each label to route each record. You sample about 500 records and label them with your large model. Then you test every observed confidence value as a routing threshold and choose the cheapest threshold that still meets your target. The same sample also shows whether the small model’s logits are correlated strongly enough with the large model’s labels for the approach to work.
Important note: the accuracy target measures agreement with the large model, not human labels, because this technique is about mimicking a large model’s answers. Measuring agreement with human labels is a different exercise.
Across eight datasets, the BARGAIN paper reports up to 86% more cost reduction than competing methods.
The follow-up Task Cascades paper adds more optimizations, such as:
- Rewriting the prompt into simpler surrogate questions a cheap model can answer.
- Reading the most relevant chunks of each document instead of the full text.
- Searching over candidate cascades to find the cheapest sequence that meets the accuracy target.
These optimizations cut costs a further 48.5% on average over BARGAIN-style cascades.
You can watch the full talk here.
Want to learn more about AI Evals? Check out our AI Evals course. It’s a live cohort with hands on exercises and office hours.