A Multi-Model Adaptation of Speculative Decoding for Classification

TL;DR

This work adapts speculative decoding from generation to classification by deploying up to three lightweight worker models to predict class labels and a single robust judge model to resolve disagreements, thereby lowering judge usage and latency. Empirical results show that out-of-the-box 3B workers can align with the judge as well as 7B workers on sentiment and similar-ticket tasks, while delivering substantial speedups up to 9x. The approach leverages non-overlapping training data across models and a confidence-threshold based majority rule to decide final labels, highlighting the practicality of scalable evaluation with open-source models. Limitations include reliance on simple majority criteria and output-format sensitivity, pointing to future work on additional agreement schemes and task-specific fine-tuning to further stabilize performance. This has practical implications for low-latency, scalable evaluation and calibration of LLMs in real-world settings.

Abstract

The current study introduces a novel adaptation of speculative decoding, repurposed from generation to classification tasks. We propose a multi-model framework employing up to three lightweight worker models and a single, more robust judge model analogous to draft models and target model, respectively, in speculative decoding. The worker models, tasked with the bulk of the computation, independently predict discrete class labels for a given input. When majority worker models agree on a label, it is accepted as the final label, optimizing efficiency by bypassing the computationally expensive judge model. In cases of disagreement, the judge model intervenes to resolve the label. This approach minimizes redundant computation, leverages the redundancy of multiple workers for confidence, and confines the judge model's role to challenging cases, offering a practical balance of efficiency and accuracy. Our analysis suggests that smaller out of the box instruction/chat finetuned worker models with 3 billion parameters (hereafter, 3B) demonstrate a level of alignment with judge models comparable to that of larger finetuned worker models with 7 billion parameters (hereafter, 7B) across both simple and higher order reasoning tasks. The top performing 3B worker model pair achieve an agreement rate of approximately 80-83% for sentiment and around 50-80% for similar ticket when compared to judge models. Additionally, 3B worker models provide a speedup ranging from 2.8x to 9x relative to the judge models, while 7B worker model combinations achieve a speedup ranging from 1.28x to 0.28x

Figures (4)

• Figure 1: Agreement of worker models with each judge for sentiment and similar ticket task.
• Figure 2: Agreement of best worker combinations with judge models for sentiment and similar ticket. The best worker pair for similar ticket is (Phi-3-mini- 4k-instruct, h2o- danube3.1-4b-cha) in 3B category and (Mistral-7B-Instruct-v0.3, Llama-3.1-8B-Instruct) in 7B category.
• Figure 3: The p95 latency of the Qwen-32B-Preview judge model and the worker model was measured using vLLM with Flash Attention 1 on a single A100 GPU with 80GB VRAM
• Figure 4: The p95 latency of the Mixtral-8x7B-Instruct judge model and the worker model was measured using vLLM with Flash Attention 1 on a single A100 with 80GB RAM.