Fuel Gauge: Estimating Chain-of-Thought Length Ahead of Time in Large Multimodal Models

1. Problem Statement

Reasoning Large Multimodal Models (LMMs) have achieved state-of-the-art performance by utilizing Chain-of-Thought (CoT) reasoning, where models decompose problems into intermediate steps. However, this approach introduces two critical inefficiencies:

• Computational Inefficiency (Memory Fragmentation): Due to the auto-regressive nature of LMMs, the final CoT length is unknown a priori. Serving frameworks must repeatedly allocate small, contiguous memory blocks for Key-Value (KV) caches as the CoT grows. This leads to severe memory fragmentation, causing allocation failures even when sufficient total memory exists.
• Sub-optimal Reasoning Quality: Without knowing the task difficulty beforehand, models often suffer from over-thinking (wasting resources on unnecessary steps) or under-thinking (stopping prematurely), leading to reduced accuracy.

Current systems lack a mechanism to predict the required CoT length before generation completes, preventing proactive resource management or dynamic intervention.

2. Methodology: Fuel Gauge

The authors propose Fuel Gauge, the first framework to estimate CoT length in real-time. The method is grounded in two core hypotheses inspired by human cognition and empirical observations:

Core Hypotheses

1. Predictability (Hypothesis I): CoT length follows a Bernoulli process and is predictable based solely on the input prompt. Empirical data shows a negative correlation between task difficulty (measured by accuracy) and CoT length; harder tasks require longer CoTs, but this relationship is consistent across models.
2. Internal "Fuel" Signal (Hypothesis II): LMMs possess an internal hidden signal analogous to biological "fuel" (e.g., ATP/adenosine levels in the brain). This signal starts at a high value (1.0) when reasoning begins and linearly decreases to zero (0.0) as the reasoning concludes.

Architecture and Algorithm

Fuel Gauge operates in two stages using a lightweight neural network (only 82k parameters):

• Stage 1: Fuel Level Estimation

  • Input: The most recent 8 hidden states ($h_{i-7:i}$) from a specific transformer layer.
  • Signal Extractor ($f_{sig}$): A 1D depth-wise convolution followed by a point-wise convolution that extracts a hidden signal vector $S_i$.
  • Fuel Estimator ($f_{fuel}$): A 2-layer MLP that maps $S_i$ to a scalar fuel level $r_i \in [0, 1]$.
  • Training: The network is trained to predict the normalized token index ($1 - i/N$) using a Smooth L1 loss, where $N$ is the ground-truth CoT length.

• Stage 2: CoT Length Prediction

  • The system collects fuel level readings ($r_0, r_1, \dots, r_i$) over time.
  • It fits a linear model ($r_i \approx k \cdot i + 1$) to these points.
  • The predicted CoT length ($\tilde{N}$) is calculated by extrapolating the line to the zero-crossing point: $\tilde{N} = -1/k$.

3. Key Contributions

1. Mathematical Characterization: The paper establishes that CoT length can be modeled as a predictable Bernoulli process dependent on task difficulty, challenging the notion that CoT length is entirely stochastic.
2. First CoT Length Predictor: Introduces Fuel Gauge, a method to extract a hidden "fuel level" signal from LMM hidden states to predict the remaining CoT length before generation ends.
3. Downstream Applications:
   • Predictive KV Cache Allocation: Uses the prediction to allocate memory in large contiguous blocks upfront, drastically reducing fragmentation.
   • CoT Length Modulation: Uses the fuel signal to guide generation via gradient-based steering (similar to classifier guidance), allowing users to force shorter or longer reasoning chains to optimize for speed or accuracy.
4. Extensive Validation: Validated across text-only, image-text, and video-text modalities on diverse benchmarks (GPQA-Diamond, MathVision, LongVideoBench).

4. Experimental Results

The authors evaluated Fuel Gauge on models like Qwen3 (4B/8B), Qwen3VL, and Intern-S1.

• Prediction Accuracy:
  • On the GPQA-Diamond benchmark, Fuel Gauge achieved a relative Mean Absolute Error (rMAE) of 0.2732 (Qwen3-8B), which is less than half the error of the best baseline ("Direct" prediction).
  • It significantly outperformed static baselines (Mean/Median) and probability-based baselines (End-of-CoT token probability), which failed to generalize across modalities.
• Memory Efficiency (Predictive KV Cache):
  • By predicting the required cache size, Fuel Gauge reduced memory allocation frequency by 13.37× on the MathVision-m benchmark (Qwen3VL-2B) compared to the standard HuggingFace (HF) allocator.
  • This effectively mitigated memory fragmentation, allowing successful inference on tasks that would otherwise fail due to fragmentation.
• CoT Modulation:
  • The method demonstrated a linear relationship between the modulation factor ($\eta$), CoT length, and model accuracy.
  • Users could reliably control the model to "think less" (shorter CoT) or "think more" (longer CoT), directly influencing accuracy in a predictable manner.
• Generalizability: The model, trained on only 200 samples from MMLU/MMMU, successfully generalized to out-of-domain tasks (e.g., video reasoning) and different modalities without retraining.
• Overhead: The computational overhead is negligible. The 82k-parameter network processes tokens at 792.7 tokens/sec (batch size 1) on an A6000 GPU, compared to the 4B model's 22.4 tokens/sec, adding minimal latency.

5. Significance

• System Optimization: Fuel Gauge solves a critical infrastructure problem in LLM serving (memory fragmentation) by enabling proactive resource allocation, which is essential for deploying long-context reasoning models in production.
• Controllability: It provides the first reliable mechanism for test-time scaling of reasoning depth using classifier guidance. This allows for dynamic trade-offs between inference cost and accuracy based on user needs.
• Theoretical Insight: The work offers a novel perspective on LLM internals, suggesting that reasoning processes have an inherent, decodable "energy" state that correlates with the completion of the task, bridging the gap between biological cognition analogies and deep learning mechanics.

In summary, Fuel Gauge transforms the unpredictable nature of Chain-of-Thought reasoning into a manageable, estimable resource, offering immediate practical benefits for system efficiency and model controllability.