Inference acceleration for large language models using "stairs" assisted greedy generation

TL;DR

The paper addresses the high inference latency of billion-parameter LLMs by introducing stairs-assisted greedy generation, which leverages a small assistant model to generate a batch of candidate tokens and a stairs-based validation to prune outputs, effectively reducing main-model iterations. This hybrid approach aims to preserve accuracy while lowering latency. Experiments on T5-large and T5-3B show substantial inference-time reductions of $17.24\%$ and $9.58\%$, respectively, with BLEU scores staying in the $75$–$100$ range, outperforming a baseline and, in some cases, the HuggingFace-assisted method. The findings suggest that adjusting the assistant-batch size can yield up to ~2x improvements in throughput, offering a practical deployment strategy and motivating future work on more prompts, tasks, and architectures.

Abstract

Large Language Models (LLMs) with billions of parameters are known for their impressive predicting capabilities but require lots of resources to run. With their massive rise in popularity, even a small reduction in required resources could have an impact on environment. On the other hand, smaller models require fewer resources but may sacrifice accuracy. In this work, we are proposing an implementation of ``stairs'' assisted greedy generation. It is a modified assisted generation methodology that makes use of a smaller model's fast generation, large model's batch prediction, and "stairs" validation in order to achieve a speed up in prediction generation. Results show between 9.58 and 17.24 percent inference time reduction compared to a stand-alone large LLM prediction in a text generation task without a loss in accuracy.

Figures (9)

• Figure 1: An abstract example of the LLM next token prediction. The token "cute" is chosen because it has the highest probability in the vocabulary. The probability is calculated based on a given input.
• Figure 2: An abstract example of the LLM next several tokens prediction. Specifically, showing how the model greedily chooses the next token in each iteration. This example took 3 iterations.
• Figure 3: An abstract example of the LLM next several tokens prediction with an assistant model emphasizing how the assistant model generates several tokens in advance. The large model keeps only those tokens that are greedily matched, and using them generates the following tokens. This example took one large model iteration.
• Figure 4: An abstract example of the LLM next several tokens prediction with "stairs" batch prediction of an assistant model emphasizing how the assistant model generates several tokens in advance and puts them into the batch. The large model predicts each of them individually. After prediction, "stairs" greedy selection is applied and only validated tokens are kept. This example took one iteration of the large model to generate the final result.
• Figure 5: "Stairs" greedy validation example that is applied after the main model predicts a given batch. The first vector is taken as a ground truth, at the same time approving tokens vertically below. Next vector's final token is checked against the vocabulary and if approved new ground truth is set. The process continues until the end of a batch or when the ground truth does not match the next vector prediction. The final ground truth is returned as a final output for a given prediction iteration.