Seq-VCR: Preventing Collapse in Intermediate Transformer Representations for Enhanced Reasoning

TL;DR

This work identifies representation collapse in intermediate Transformer layers as a key bottleneck for multi-step arithmetic reasoning and introduces Sequential Variance-Covariance Regularization (Seq-VCR) to preserve representation diversity. By combining Seq-VCR with dummy pause tokens that mimic chain-of-thought without supervision, the method achieves strong reasoning performance, notably reaching 99.5% exact match on 5×5 multiplication and surpassing GPT-4 with five-shot CoT in some settings. Seq-VCR promotes high-variance, low-covariance intermediate representations, evidenced by higher layer entropy and improved learning dynamics, while also reducing inference time relative to explicit CoT. The approach yields significant gains on Arithmetic Expressions and LIS tasks, highlighting its potential as a robust,CoT-free mechanism to enhance transformer reasoning. These results underscore the importance of preventing intermediate-collapse and suggest broader applicability to diverse sequential reasoning problems.

Abstract

Decoder-only Transformers often struggle with complex reasoning tasks, particularly arithmetic reasoning requiring multiple sequential operations. In this work, we identify representation collapse in the model's intermediate layers as a key factor limiting their reasoning capabilities. To address this, we propose Sequential Variance-Covariance Regularization (Seq-VCR), which enhances the entropy of intermediate representations and prevents collapse. Combined with dummy pause tokens as substitutes for chain-of-thought (CoT) tokens, our method significantly improves performance in arithmetic reasoning problems. In the challenging $5 \times 5$ integer multiplication task, our approach achieves $99.5\%$ exact match accuracy, outperforming models of the same size (which yield $0\%$ accuracy) and GPT-4 with five-shot CoT prompting ($44\%$). We also demonstrate superior results on arithmetic expression and longest increasing subsequence (LIS) datasets. Our findings highlight the importance of preventing intermediate layer representation collapse to enhance the reasoning capabilities of Transformers and show that Seq-VCR offers an effective solution without requiring explicit CoT supervision.

Figures (15)

• Figure 1: Position-wise number of operations needed for 5x5 digits integer multiplication task. Middle tokens in the output sequence need more operations than the peripheral ones, making their prediction much harder (as shown in Figure \ref{['fig:position_acc']}). Example of 12345 x 67890 is shown here.
• Figure 2: Representation collapse across layers during (a) training (or fine-tuning(b)) for two datasets. The x-axis represents the layer ID, while the y-axis shows the degree of collapse as measured by representation Matrix-Entropy. The results highlight how intermediate layers (shown by the decline in Entropy) experience representation collapse for Pretrained GPT-2 Small on 5 × 5 digit Multiplication (b) and during Vanilla training or fine-tuning for both datasets, indicating potential bottlenecks in information flow or feature learning. Tools like Pausegoyal2023think token-based tuning can't fix it, but our proposed regularization Seq-VCR can improve collapse.
• Figure 3: Illustrations of Input, Output and CoT on the Arithmetic, LIS and mutliplication datasets
• Figure 4: Layer-wise entropy distributions for different configurations on the $5 \times 5$ multiplication task. Seq-VCR and Seq-VCR+Pause maintain higher entropy across layers, indicating greater representation diversity.
• Figure 5: Learning dynamics (curves for next token prediction loss) illustrating the phase transition observed across datasets when applying our regularization methods. The x-axis represents training epochs, and the y-axis denotes the model's loss. The phase transition is characterized by a sharp reduction in loss, marking a distinct shift in the learning regime when using Seq-VCR and Seq-VCR + Pause, compared to the gradual decline or saturating curves in other configurations.