🤖 AI Summary
This study addresses the instability in long-sequence generation caused by error accumulation during multi-step reasoning. It establishes the first theoretical stability bounds based on the spectral norm of Jacobian products, proving that Transformers can converge to zero-loss solutions. Guided by this analysis, the work proposes chain-of-thought compression, quantization-aware training, and a novel regularization technique to suppress error propagation. The effectiveness of these methods is validated on graph algorithm and symbolic state tracking tasks, achieving an average benchmark performance improvement of 3.5% and up to 8.2% in long-input scenarios. Furthermore, stability metrics are reduced by three- to eight-fold, effectively mitigating the exponential growth of errors.
📝 Abstract
We consider the stability of multi-step reasoning processes, which have extensive applications in language models, including chain-of-thought and algorithmic reasoning. While longer sequences of reasoning can improve a model's generation capability at test time, the errors due to intermediate reasoning steps can accumulate in autoregressive generation, and thus grow substantially at the end. In this paper, we ask: What are the key factors determining the stability of multi-step reasoning? First, we show an inference error bound governed by the product of spectral norms of the Jacobians taken through the input space across generation steps. This product can be viewed as an error amplification factor, which could scale exponentially with the number of reasoning steps, serving as a quantitative measure of reasoning stability. Second, we analyze this measure in transformer models trained to predict simple tasks like linear and quadratic functions. We theoretically prove that the transformer model converges to a solution where the stability measure decays, thus yielding nearly zero inference loss over (arbitrarily) long steps. Finally, the stability analysis leads to several algorithmic implications for controlling the stability, through (i) chain-of-thought length compression that reduces the sensitivity of each step, and (ii) quantization-aware training that regularizes the input Jacobian norms. We validate the proposed algorithms by fine-tuning language models on graph-algorithmic reasoning tasks and symbolic state-tracking tasks. Across seven evaluations, our algorithms improve over baseline comparisons by 3.5% on average, and by 8.2% for longer-length inputs. Ablation analysis validates that the stability measure is drastically reduced by 3-8$\times$, confirming the regularization effect on the spectral norms of the (input space) Jacobians.