Fast Chain-of-Thought: A Glance of Future from Parallel Decoding Leads to Answers Faster

TL;DR

This work targets the bottleneck of autoregressive decoding in Chain-of-Thought tasks by introducing FastCoT, a model-agnostic framework that provides a glimpse of the future through parallel Jacobi decoding while maintaining exact tokens through a lossless decoding path. It introduces an Approximate Tokens Buffer to manage exact and approximate tokens and enables parallel forward passes, reducing inference time by up to ~20% with negligible accuracy loss across multiple datasets and model scales. Key contributions include the integration of Jacobi decoding with a practical batch-friendly implementation, detailed analyses of time costs and context-window effects, and demonstrations of robustness across tasks and models. The approach offers a pathway to faster, high-throughput CoT reasoning in real-world settings without additional training or model modifications, while highlighting design considerations such as context-window size and KV-cache handling.

Abstract

In this work, we propose FastCoT, a model-agnostic framework based on parallel decoding without any further training of an auxiliary model or modification to the LLM itself. FastCoT uses a size-varying context window whose size changes with position to conduct parallel decoding and auto-regressive decoding simultaneously, thus fully utilizing GPU computation resources. In FastCoT, the parallel decoding part provides the LLM with a quick glance of the future composed of approximate tokens, which could lead to faster answers compared to regular autoregressive decoding used by causal transformers. We also provide an implementation of parallel decoding within LLM, which supports KV-cache generation and batch processing. Through extensive experiments, we demonstrate that FastCoT saves inference time by nearly 20% with only a negligible performance drop compared to the regular approach. Additionally, we show that the context window size exhibits considerable robustness for different tasks.

Figures (8)

• Figure 1: Example of how the glimpse of future work in CoT reasoning task. In the CoT reasoning task, LLM is required to completely generate the complete rationale and finally get the answer based on the rationale. FastCoT argues that partial generation of the complete rationale with a glimpse of future is enough. The example in this figure comes from our experiment results.
• Figure 2: Time cost with context window size $c$. The time overhead does not increase significantly as the window size increases. Almost no change in time consumption from autoregressive decoding (size=0) to Jacobi decoding with a window size of 38. We apply the Llama2-13B model with a single Nvidia A100 GPU.
• Figure 3: Overview of FastCoT. The autoregressive decode get another one exact token, further future’s approximate token is revealed by Jacobi decoding.
• Figure 4: Overview of a complete iteration of FastCoT.
• Figure 5: Reasonable overlap ratio would lead to a saturated accuracy.