Cascade-Aware Training of Language Models

TL;DR

This paper introduces cascade-aware training (CAT), a novel method for optimizing cascaded language models by considering inference-time interactions, leading to better quality-cost tradeoffs across multiple benchmark tasks.

Contribution

We propose cascade-aware training (CAT), a new approach that trains small language models with awareness of their role in cascades, improving overall performance and efficiency.

Findings

Improves quality-cost tradeoff in cascaded LMs

Demonstrates effectiveness on 60+ LM tasks

Enhances inference efficiency in practical settings

Abstract

Reducing serving cost and latency is a fundamental concern for the deployment of language models (LMs) in business applications. To address this, cascades of LMs offer an effective solution that conditionally employ smaller models for simpler queries. Cascaded systems are typically built with independently trained models, neglecting the advantages of considering inference-time interactions of the cascaded LMs during training. In this paper, we present cascade-aware training(CAT), an approach to optimizing the overall quality-cost performance tradeoff of a cascade of LMs. We achieve inference-time benefits by training the small LM with awareness of its place in a cascade and downstream capabilities. We demonstrate the value of the proposed method with over 60 LM tasks of the SuperGLUE, WMT22, and FLAN2021 datasets.