Table of Contents
Fetching ...

Accelerating Chain-of-Thought Reasoning: When Goal-Gradient Importance Meets Dynamic Skipping

Ren Zhuang, Ben Wang, Shuifa Sun

TL;DR

This work tackles the inefficiency of Chain-of-Thought reasoning in large language models by introducing Adaptive GoGI-Skip, a framework that jointly learns a goal-oriented token importance metric (GoGI) and a dynamic skipping mechanism (ADS). GoGI measures a token’s influence on the final answer loss via intermediate-gradient signals, while ADS uses Entropy-Driven Rate and Adaptive N-Constraints to regulate pruning and maintain coherence. Trained with compressed MATH data and evaluated across diverse benchmarks (AIME, GPQA, GSM8K) and model families (Gemma3, Qwen2.5), the method achieves over 45% average CoT token reduction and 1.6–2.0x inference speedups with minimal accuracy loss, outperforming static baselines. The approach demonstrates substantial cross-domain generalization and highlights the value of combining goal-oriented importance with uncertainty-aware, coherence-preserving adaptation for efficient, high-fidelity reasoning in LLMs.

Abstract

Large Language Models leverage Chain-of-Thought (CoT) prompting for complex tasks, but their reasoning traces are often excessively verbose and inefficient, leading to significant computational costs and latency. Current CoT compression techniques typically rely on generic importance metrics and static compression rates, which may inadvertently remove functionally critical tokens or fail to adapt to varying reasoning complexity. To overcome these limitations, we propose Adaptive GoGI-Skip, a novel framework learning dynamic CoT compression via supervised fine-tuning. This approach introduces two synergistic innovations: (1) Goal-Gradient Importance (GoGI), a novel metric accurately identifying functionally relevant tokens by measuring the gradient influence of their intermediate representations on the final answer loss, and (2) Adaptive Dynamic Skipping (ADS), a mechanism dynamically regulating the compression rate based on runtime model uncertainty while ensuring local coherence through an adaptive N-token constraint. To our knowledge, this is the first work unifying a goal-oriented, gradient-based importance metric with dynamic, uncertainty-aware skipping for CoT compression. Trained on compressed MATH data, Adaptive GoGI-Skip demonstrates strong cross-domain generalization across diverse reasoning benchmarks including AIME, GPQA, and GSM8K. It achieves substantial efficiency gains - reducing CoT token counts by over 45% on average and delivering 1.6-2.0 times inference speedups - while maintaining high reasoning accuracy. Notably, it significantly outperforms existing baselines by preserving accuracy even at high effective compression rates, advancing the state of the art in the CoT reasoning efficiency-accuracy trade-off.

Accelerating Chain-of-Thought Reasoning: When Goal-Gradient Importance Meets Dynamic Skipping

TL;DR

This work tackles the inefficiency of Chain-of-Thought reasoning in large language models by introducing Adaptive GoGI-Skip, a framework that jointly learns a goal-oriented token importance metric (GoGI) and a dynamic skipping mechanism (ADS). GoGI measures a token’s influence on the final answer loss via intermediate-gradient signals, while ADS uses Entropy-Driven Rate and Adaptive N-Constraints to regulate pruning and maintain coherence. Trained with compressed MATH data and evaluated across diverse benchmarks (AIME, GPQA, GSM8K) and model families (Gemma3, Qwen2.5), the method achieves over 45% average CoT token reduction and 1.6–2.0x inference speedups with minimal accuracy loss, outperforming static baselines. The approach demonstrates substantial cross-domain generalization and highlights the value of combining goal-oriented importance with uncertainty-aware, coherence-preserving adaptation for efficient, high-fidelity reasoning in LLMs.

Abstract

Large Language Models leverage Chain-of-Thought (CoT) prompting for complex tasks, but their reasoning traces are often excessively verbose and inefficient, leading to significant computational costs and latency. Current CoT compression techniques typically rely on generic importance metrics and static compression rates, which may inadvertently remove functionally critical tokens or fail to adapt to varying reasoning complexity. To overcome these limitations, we propose Adaptive GoGI-Skip, a novel framework learning dynamic CoT compression via supervised fine-tuning. This approach introduces two synergistic innovations: (1) Goal-Gradient Importance (GoGI), a novel metric accurately identifying functionally relevant tokens by measuring the gradient influence of their intermediate representations on the final answer loss, and (2) Adaptive Dynamic Skipping (ADS), a mechanism dynamically regulating the compression rate based on runtime model uncertainty while ensuring local coherence through an adaptive N-token constraint. To our knowledge, this is the first work unifying a goal-oriented, gradient-based importance metric with dynamic, uncertainty-aware skipping for CoT compression. Trained on compressed MATH data, Adaptive GoGI-Skip demonstrates strong cross-domain generalization across diverse reasoning benchmarks including AIME, GPQA, and GSM8K. It achieves substantial efficiency gains - reducing CoT token counts by over 45% on average and delivering 1.6-2.0 times inference speedups - while maintaining high reasoning accuracy. Notably, it significantly outperforms existing baselines by preserving accuracy even at high effective compression rates, advancing the state of the art in the CoT reasoning efficiency-accuracy trade-off.
Paper Structure (92 sections, 11 equations, 8 figures, 8 tables, 1 algorithm)

This paper contains 92 sections, 11 equations, 8 figures, 8 tables, 1 algorithm.

Figures (8)

  • Figure 1: Building an accelerated language model with Adaptive GoGI-Skip framework. (a) Calculating GoGI by measuring the gradient influence of each token's intermediate representation on the final answer loss for an original CoT. (b) Applying ADS, informed by EDR regulation and ANC, to determine which tokens to prune based on both importance and local context. (c) This dynamically compressed CoT data then fine-tunes a base LLM, enabling efficient inference without requiring additional parameters at runtime.
  • Figure 2: GoGI score characteristics and dynamic behavior of Adaptive GoGI-Skip.Left: The highly skewed distribution of GoGI scores from 7,472 MATH training samples. Right: Illustration of the dynamic token retention rate's multi-faceted response to runtime entropy, original CoT length, and overall GoGI score distribution. The color gradient indicates the average GoGI score of the tokens retained by the dynamic mechanism.
  • Figure 3: Properties of GoGI and the ANC Mechanism.Left: A hexbin plot of per-token GoGI scores versus predictive entropy illustrates their near-orthogonality, suggesting GoGI captures a distinct dimension of importance. Right: The ANC mechanism dynamically adjusts the maximum number of consecutive token skips ($N_t$) based on local contextual stability (derived from windowed predictive entropy), thereby ensuring reasoning coherence.
  • Figure 4: Analysis supporting GoGI target layer selection for Gemma3-4B-it.Left: Layer-wise gradient contribution to logits. Right: GoGI statistics across layers. While high max GoGI occurs near L23 and L28, L23 shows stronger overall contribution, which informs our choice.
  • Figure 5: Distribution of normalized gradient contributions from each layer to the final logits across all evaluated models.Top: Gemma3-1B-it and Gemma3-12B-it. Bottom: Qwen2.5-1B-it and Qwen2.5-7B-it.
  • ...and 3 more figures