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

## Abstract

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 1 1 1 https://github.com/rarefin/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 × 5 5 5 5\times 5 5 × 5 integer multiplication task, our approach achieves 99.5 % percent 99.5 99.5\% 99.5 % exact match accuracy, outperforming models of the same size (which yield 0 % percent 0 0\% 0 % accuracy) and GPT-4 with five-shot CoT prompting ( 44 % percent 44 44\% 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.

## 1 Introduction

Figure: 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 6). Example of 12345 x 67890 is shown here.
Refer to caption: x1.png

Large Language Models (LLMs) based on Transformer architectures have achieved remarkable success across a wide range of tasks, positioning them as foundational models in artificial intelligence (Bommasani et al., ). Despite their impressive capabilities, LLMs often struggle with tasks requiring complex reasoning, particularly arithmetic reasoning that necessitates multiple sequential operations (Bubeck et al., ). These challenges are attributed to the models’ limitations in handling tasks that involve deep cognitive abilities and multi-step reasoning processes.

One of the key obstacles is the representation collapse in the intermediate layers of Transformer models. Representation collapse occurs when internal representation diversity diminishes, leading to less informative features and hindering the model’s ability to solve complex tasks (Jing et al., ). For arithmetic reasoning, transformers struggle with successive carryovers and storing intermediate results Qiu et al. (), which are essential for solving complex sub-tasks, requiring more computations (Figure 1). We hypothesize that representation collapse prevents the model from effectively performing sub-tasks by calculating successive carryovers and storing intermediate results, which are essential for accurate prediction.

To address this limitation, we introduce Sequential Variance-Covariance Regularization (Seq-VCR), a regularization technique designed to enhance the entropy of intermediate representations and prevent representation collapse. By increasing the diversity of representations within the model’s layers, Seq-VCR enables the Transformer to maintain richer and more informative features throughout the computation process.

Furthermore, we incorporate dummy pause tokens as substitutes for chain-of-thought (CoT) tokens. While CoT prompting has been shown to improve reasoning by breaking down tasks into intermediate steps (Wei et al., ), it often requires explicit supervision and can be computationally expensive. Our approach leverages pause tokens to simulate the effect of CoT without the need for explicit intermediate reasoning steps.

We validate our method on challenging arithmetic reasoning tasks. Notably, on the $5\times 5$ integer multiplication task, our approach achieves $99.5\%$ exact match accuracy, surpassing models of the same size (which yield $0\%$ accuracy) and GPT-4 with five-shot CoT prompting ($44\%$). We also demonstrate significant improvements on arithmetic expression and longest increasing subsequence (LIS) datasets.

Our contributions are summarized as follows:

- •
We identify representation collapse in intermediate layers as a key limitation affecting the reasoning capabilities of Transformer models.
- •
We propose Seq-VCR, a regularization technique that enhances the diversity of intermediate representations and prevents collapse.
- •
We demonstrate that combining Seq-VCR with pause tokens enables models to solve complex arithmetic reasoning tasks without explicit CoT supervision.
- •
We provide extensive experimental results showing significant improvements over baseline models and state-of-the-art LLMs like GPT-4.

## 2 Literature Review

##### Transformer Reasoning

Research in large language models (LLMs) has advanced significantly. Models like BERT (Devlin et al., ) and GPT (Radford & Narasimhan, ) demonstrated improvements in natural language understanding. However, challenges remain in tasks that require deep cognitive abilities. Integrating external knowledge has shown promise in improving reasoning capabilities (Bosselut et al., ). Recent advances include techniques such as chain-of-thought (CoT) prompting, which allow models such as GPT-3 (Brown et al., ) to perform complex reasoning by breaking tasks into intermediate steps (Wei et al., ) with human supervision, called process supervision (Lightman et al., ). Wang et al. () shows that even incorrect but coherent intermediate steps can improve reasoning performance. Jin et al. () found increasing the length of the reasoning steps in the prompts, even without adding new information to the prompt, improves the reasoning abilities of LLMs.

Wang et al. (), investigated the addition of dummy tokens such as periods or hash symbols to inputs at inference time. They found that this simple modification could improve performance in arithmetic reasoning tasks. Some works show that these dummy tokens, commonly referred to as filler tokens, do not extend transformers’ abilities beyond TC^0 circuit complexity, but still significantly enhance problem-solving within this class (Merrill & Sabharwal, ; Strobl et al., ).
Building on this idea, Goyal et al. () introduced a more comprehensive framework called ”pause-training.” Their approach involves incorporating learnable $<pause>$ tokens during both pre-training and fine-tuning of language models. While there are differences between the pause tokens and filler tokens approach, they share the goal of providing additional computation time for transformers which extends the expressive power of transformers within the TC^0 class without changing the fundamental limitations of the model architecture. Chain-of-thought prompting can potentially elevate transformers beyond the TC^0 complexity class; However, it may not be necessary for all problems.
Furthermore, filler tokens have been demonstrated to enhance the expressivity of transformers (Pfau et al., ), even in cases where traditional transformers, lacking chain-of-thought mechanisms, exhibit insufficient expressive power (Sanford et al., ).

##### Representation Collapse

Representation collapse in representation learning degrades model performance by making features indistinguishable. It often occurs in unsupervised and self-supervised learning tasks due to the lack of labeled data, leading to trivial solutions and a low-rank feature space.
Research by Jing et al. () highlights that representation collapse is due to poor optimization, loss function design, and improper regularization. Recent studies (Grill et al., ) suggest the use of contrastive learning, such as SimCLR, to mitigate this by contrasting positive and negative samples. Additionally, Zbontar et al. () introduced Barlow Twins to minimize redundancy between learned representations. A further advancement is VICReg (Bardes et al., ; Zhu et al., ) which employs variance, invariance, and covariance regularization to prevent collapse, ensuring diverse and non-trivial representations. While these researches focuse on images, it is less explored in language models. Recently, Barbero et al. () studied last-layer representations of the next tokens and found representation collapse in pre-trained models.

## 3 Methodology

### 3.1 Preliminaries

Consider a finite set of token vocabulary denoted by $V=(1\ldots v)$, while $x,y\in V$ constitutes a sequence of these tokens representing the input $\mathbf{x}=[x_{1},\ldots,x_{K}]$ and the output $\mathbf{y}=[y_{1},\ldots,y_{T}]$, where $K$ and $T$ indicate the respective sequence lengths of the input and output.

We covert the input $\mathbf{x}$ and output sequence $\mathbf{y}$ in an embedded
sequence:

$E_{emb}=[\tilde{x}_{1},\ldots,\tilde{x}_{K},\tilde{y}_{1},\ldots,\tilde{y}_{T}
]\in\mathbb{R}^{T+K\times d}$ and $E_{pos}=[p_{1},p_{2},\ldots]\in\mathbb{R}^{T+K\times d}$ of dimension $d$.

We can get the language model layer-wise as:
$f=(f_{cls}\circ f_{L}\circ f_{L-1}\ldots\circ f_{0})$

where $f_{0}=[\tilde{x}_{1}+p_{1},\ldots,\tilde{x}_{K}+p_{K},\tilde{y}_{1}+p_{K+1},
\ldots,\tilde{y}_{T}+p_{K+T}]$ and $f_{cls}\in\mathbb{R}^{|V|\times d}$ is the final linear classification layer.
$f_{l}$ is the self-attention layer with MLP for layer $l$ (Vaswani et al., ).

### 3.2 Next-Token Prediction Loss

Large language models like *GPT* (Radford et al., ) (Generative Pre-trained Transformer) are trained to auto-regressively predict the next token in a sequence. This minimizes the difference between predicted and actual tokens for coherent and contextually appropriate text. A loss function like cross-entropy for next-token prediction is given by:

$$ $\mathcal{L}_{next}(y,f(x))=-\sum_{t=1}^{T}y_{t}^{\top}\log(f(x,y_{<t})),$ (1) $$

where $y_{t}$ is the one-hot encoded true token at time step $t$ and input tokens $x$, modeled by a language model, $f$.

### 3.3 Representation Collapse in Transformers

Figure: (a) Training minGPT from scratch on *Arithmetic Expression*
Refer to caption: x2.png

Representation collapse in Transformers manifests as a reduction in the diversity of internal representations across layers, particularly in tasks requiring complex reasoning. This collapse hinders the model’s ability to capture and process intricate patterns necessary for tasks like multi-digit multiplication.

To illustrate this, we analyze the layer-wise representation entropy of a GPT-2 model fine-tuned on the $5\times 5$ digit multiplication task. As shown in Figure 2(b), there is a significant drop in entropy in the intermediate layers, indicating a collapse in representation diversity. We hypothesize that this representation collapse is the reason why it is very hard for the model to calculate multiplication, carry, and store them, thus resulting in poorer performance, as shown in Figure 6 for the output positions, which require more computation (Figure 1).

For computing the representation collapse of a token sequence, we utilise the matrix entropy of the representation vectors in different layers of the transformer model. We investigate $\alpha$-order matrix-based entropy (Giraldo et al., ; Skean et al., ) as a practical alternative to Rényi’s $\alpha$-order entropy (Rényi, ). This metric uses a similarity kernel $\kappa$ on layer-wise representations of token sequences from a distribution, without assuming the true nature of the distribution. The kernel selection $\kappa$ is flexible, including Gaussian, linear, or Laplacian kernels. We choose the linear kernel $\kappa(a,b)=ab^{T}$ due to the linear representation hypothesis (Park et al., ), suggesting that language model representations encode high-level concepts in linearly separable manifolds.

Aussuming $Z^{(l)}=f_{l}\in\mathbb{R}^{T\times d}$ is the $l^{th}$ representations of the $T$ tokens in the sequence of dimension $d$, we can first construct the Gram matrix $\mathbf{K}=\kappa\left(Z^{(l)},Z^{(l)}\right)=Z^{(l)}Z^{(l)^{T}}\in\mathbb{R}^
{\textrm{T}\times\textrm{T}}$ consisting of all pairwise evaluations of the $T$ token representations in $Z^{(l)}$. Then the matrix-based entropy of order $\alpha>0$ is defined as:

$$ $S_{\alpha}\left(Z^{l}\right)=\frac{1}{1-\alpha}\log{\left[\sum_{i=1}^{T}\left( \frac{\lambda_{i}(\mathbf{K})}{\textrm{tr}(\mathbf{K})}\right)^{\alpha}\right]}$ (2) $$

Equation 2 is the $\alpha$-order Rényi entropy of the Gram matrix eigenvalues(^2^22The non-zero eigenvalues of the Gram matrix $Z^{(l)}Z^{(l)^{T}}$ are equivalent to those of the covariance matrix $Z^{(l)^{T}}Z^{(l)}$. Using the covariance matrix instead of the Gram matrix in Equation 2 makes no difference and is more computationally efficient if $d<T$.). Each eigenvalue is divided by $\textrm{tr}(\mathbf{K})$ and raised to the $\alpha$ power so that the eigenvalues of $\mathbf{K}$ sum to $1$, as $\textrm{tr}(\cdot)=\sum_{i=1}^{n}\lambda_{i}(\cdot)$), allowing them to be treated as a probability distribution. Each eigenvalue of $\mathbf{K}$ represents the variance of the sequence token representations in a specific direction of the principal component Scholkopf & Smola (). Low entropy indicates a heavy-tailed distribution in which few components dominate the variance in $Z^{(l)}$. At maximum entropy, the eigenvalues form a uniform distribution, spreading token representations equally in all directions. Depending on $\alpha$, various special cases of matrix-based entropy can be recovered. Specifically, when $\lim_{\alpha\rightarrow 1}$ it equals Shannon entropy, also known as Von Neumann entropy in quantum information theory Bach (); Boes et al. ()).

### 3.4 Sequential Variance Covariance Regularization (Seq-VCR)

To prevent representation collapse, we propose Seq-VCR, a regularization technique that encourages the model to maintain high variance and low covariance in its intermediate representations. Seq-VCR is inspired by techniques used in self-supervised learning for image representations (Bardes et al., ) and is adapted here for sequential data in Transformer models.

Seq-VCR is applied to the final output of the model ($X=f_{cls}$). Since it requires a calculation of the covariance matrix and a pre-trained model such as GPT-2, the dimension can be $f_{cls}\mathbb{R}^{5000\times d}$, such dimensional calculation is often difficult. So we use a linear projection layer of the representation layer $f_{proj}$ which projects the representation of the layer $l$ into a smaller embedding space such as $X=f_{proj}(f_{l})R^{T\times d}$, $f_{proj}\in\mathbb{R}^{2048\times d}$.

Now we can assume that we have $N$ samples in the batch, so $N$ number of $X$. We can calculate the covariance matrix across the batch dimension, $C\in R^{T\times d\times d}$.

$$ $L_{\text{Seq-VCR}}=\frac{1}{T\times d}\sum_{i=1}^{T}\sum_{k=1}^{d}\left( \lambda_{1}\underbrace{\text{max}(0,1-\sqrt{\mathbf{C}_{i,k,k}+\eta})}_{{ \color[rgb]{1,0,0}\text{Variance Term}}}+\lambda_{2}\underbrace{\sum_{k\neq \hat{k}}(\mathbf{C}_{i,k,\hat{k}})^{2}}_{{\color[rgb]{0,0,1}\text{Covariance Term}}}\right)$ (3) $$

where $\lambda_{1}$ and $\lambda_{2}$ are the coefficients of the variance and covariance regularization terms. $\eta$ is a small constant, used for numerical stability.

The Variance Term encourages unit variance in each dimension, while the Covariance Term penalizes covariance between different dimensions, promoting decorrelation and diversity in representations.

### 3.5 Incorporating Pause Tokens

We further enhance the model’s reasoning capacity by introducing dummy pause tokens, which act as placeholders for intermediate computation steps. Unlike explicit CoT tokens, pause tokens do not require supervision and serve to allocate additional computational resources within the model (Goyal et al., ).

### 3.6 Training Objective

The overall training objective combines the standard next-token prediction loss with the Seq-VCR regularization.
Our final loss $\mathbf{L}$ is:

$$ $\mathbf{L}=L_{next}+L_{\text{Seq-VCR}}$ (4) $$

This encourages the model to not only predict the next token accurately but also maintain diverse and informative intermediate representations.

## 4 Experiments

### 4.1 Data

Figure: (a) Arithmetic Expression
Refer to caption: extracted/5970711/figures/arithmetic_ascii.png

We conduct experiments on three tasks: we first consider the multi-digit multiplication task from the BIG-bench benchmark (Srivastava et al., ), which is the most challenging
among arithmetic tasks (Yang et al., ). In particular, we use the four-digit (4 × 4) and five-digit (5 × 5) multiplication problems, since these two tasks prove very challenging to solve under no CoT, utilizing the training data generated by Deng et al. ().

Next, we focus on the Arithmetic Expressions Feng et al. () dataset. This task focuses on evaluating arithmetic expressions. The input sequence for this task is a sequence consisting of numbers, mathematical operators such as addition (+), subtraction (-), multiplication $(\times)$, division $(\div)$ and brackets, followed by the equal sign. The goal of this task is to calculate the input arithmetic expression and generate the correct result.
This task is naturally well-suited for a Chain-of-Thought (CoT) method, where each step does part of the calculation, slowly solving one operation at a time while keeping the other parts the same.

Finally, we also conduct experiments in a more general setting outside arithmetic tasks, called Dynamic Programming (DP). Dynamic Programming is a framework for solving decision-making problems. The core idea in this setting is to break down a complex task into a series of smaller subtasks that can be solved sequentially. We chose a very popular problem in this setting called the Longest Increasing Sub-sequence (LIS) as described in the Introduction to Algorithms book (Cormen et al., ). For this task, the goal is to find the length of the longest increasing subsequence of a
given integer sequence. We generate datasets with different input sequence lengths ranging from {50, 80, 100} as shown in Feng et al. (). Moreover, all input sequences, and answers in LIS are bounded-range integers and can therefore be tokenized (similar to the multi-digit multiplication task). The CoT tokens for this task consist of the dynamic programming array plus the final answer. Figure 3 shows a working example for a sample problem in all the datasets considered.

Moreover, Table 3 in the Appendix provides an overview of different tasks across the datasets i.e., Multiplication, Arithmetic Expression and LIS, along with their respective input, chain-of-thought (CoT), and output token details.

### 4.2 Models and Training Configurations

We conduct experiments using two models:

##### GPT-2 Small Fine-Tuning

We fine-tune a pre-trained GPT-2 Small model on the multi-digit multiplication tasks. Fine-tuning is performed for 40 epochs with a learning rate of $5\times 10^{-4}$ and a batch size of 32.

##### minGPT Training from Scratch

We train a minGPT model from scratch on the Arithmetic Expressions and LIS datasets. Training is conducted for 100 epochs with a learning rate of $1\times 10^{-4}$ and a batch size of 128.

We compare five configurations:

- •
Vanilla: Standard training without CoT or pause tokens.
- •
With CoT: Training with explicit CoT supervision.
- •
Pause: Inserting pause tokens in the input sequence.
- •
Seq-VCR: Applying Seq-VCR regularisation.
- •
Seq-VCR + Pause: Combining Seq-VCR with pause tokens.

### 4.3 Results

#### 4.3.1 Representation Diversity

Figure: 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.
Refer to caption: x3.png

Seq-VCR significantly enhances representation diversity, as evidenced by the layer-wise distribution of representation entropy across different configurations. In Figure 4, we observe that configurations employing our regularization technique, particularly Seq-VCR and Seq-VCR+Pause, exhibit higher entropy values in intermediate layers compared to the other configurations (Vanilla, and Pause). This increased diversity in representation entropy suggests that the model is able to capture a wider range of features and maintain distinct representations. The histograms also reveal that while other configurations show peaks at lower entropy values, indicating potential representation collapse and redundancy in the intermediate representations, our method fosters a more uniform distribution across the layers. This uniformity suggests a richer feature space being utilized by the model, enabling it to generalize better to unseen data, thus improving performance on downstream tasks.

### 4.4 Learning Dynamics

Figure: (a) Fine-tuning GPT-2 Small on *$5\times 5$ digit Multiplication*
Refer to caption: x4.png

Our regularization method, Seq-VCR or Seq-VCR+Pause cause a significant phase transition (improving the next token prediction loss) in the performance of the model across datasets compared to other configurations such as Vanilla, and Pause, which do not induce this phase transition as shown in Figure 5. Specifically, this phase transition marks a distinct shift in the model’s ability to capture and generalize representations during training or finetuning.

Without regularization, learning is difficult and training loss saturates. Seq-VCR improves the learning curve. Adding Seq-VCR with $2$ Pause tokens causes a sharp phase transition (Figure 5 (a)) and solves the $5\times 5$ digit multiplication task. Figure 5 (b) shows similar improvements in arithmetic reasoning task that is learned from scratch.

### 4.5 Results on 4 × 4 4 4 4\times 4 4 × 4 and 5 × 5 5 5 5\times 5 5 × 5 digits Multiplication Tasks

Table 1 shows the results on $4\times 4$ and $5\times 5$ digits using a fine-tuned GPT-2 Small. Compared to no CoT configurations like (Vanilla, Pause), our method enables solving tasks previously not solvable without explicit CoT. We can see the effectiveness of the Seq-VCR on $4\times 4$ digit multiplication task where without pause token it improves performance from $25\%$ to $52\%$ and with $2$ pause tokens Seq-VCR + Pause achieves accuracy $99.2\%$, which is close to With CoT performance and significantly better than other configuration without CoT. It even outperforms the 5-shot prompt(with or without CoT) performance of GPT-3.5 and GPT-4.0.

**Table 1: Accuracy (exact match) on $4\times 4$ and $5\times 5$ digits Multiplication Tasks. GPT-3.5 and GPT-4 results are taken from Deng et al. () which are produced by 5-shot prompt**
| Model | Configuration | 4x4 Mult | 5x5 Mult |
| --- | --- | --- | --- |
| GPT-3.5 | With CoT | 0.43 | 0.05 |
| No CoT | 0.02 | 0.00 |  |
| GPT-4 | With CoT | 0.77 | 0.44 |
| No CoT | 0.04 | 0.00 |  |
| GPT-2 Small | With CoT | 1.0 | 1.0 |
| Vanilla | 0.25 | 0.0 |  |
| Pause | 0.28 | 0.0 |  |
| Seq-VCR | 0.52 | 0.0 |  |
| Seq-VCR + Pause | 0.992 | 0.995 |  |

We also see from Table 1 that GPT-4 with CoT 5-shot prompts achieves $44.0\%$ for $5\times 5$ digit multiplication, while Vanilla, Seq-VCR, and Pause achieve $0\%$ without CoT. Because, all configurations struggle with predicting middle tokens in the output sequence (shown in Figure 6). This happens because those middle tokens require more computations to calculate successive carry and store them, since these tokens require more computation (as shown in Figure 1), thus making them difficult sub-task than predicting the peripheral tokens which requires less computes, thus easier to solve. However, Seq-VCR+Pause with $2$ Pause tokens solves the $5\times 5$ digit multiplication task, previously unsolved without CoT tokens. Figure 2(b) shows our regularization Seq-VCR increases intermediate layer entropy, allowing more exploration in the representation space. We hypothesize that along with better representation by Seq-VCR, $2$ Pause tokens increase the computational ability, allowing the model to carry the necessary multiplication values and solve the task.

Figure: Figure 6: Position-wise accuracy of Multiplication of 5x5 digits with different configuration of GPT-2 Small. We see models fail on the middle tokens they require more compute.

### 4.6 Results on Arithmetic Expression and Longest Increasing Sub-sequence (LIS) Datasets

Figure: (a) Test accuracy on Arithmetic Expressions Dataset
Refer to caption: extracted/5970711/figures/arithmetic_accuracy_varying_operators_with_cot.png

Figure: (a) Test accuracy on 6 operator Arithmetic Expression
Refer to caption: extracted/5970711/figures/arithmetic_6_performance_methods.png

Figure 7 illustrates how the accuracy varies across different methods (represented by different bar colors in the plot) and with changing levels of task complexity, such as the number of operators in case of the Arithmetic Expressions Dataset or input sequence length in case of the LIS Dataset.

We can observe that as the task complexity is low i.e in case of 4 operators in the arithmetic expression and input sequence length of 50 in the LIS dataset, all the methods perform quite well. But, as the number of operators or sequence length increases, the tasks become more complex which shows the need for models other than the vanilla baselines.

Seq-VCR demonstrates substantial improvement over vanilla models and achieved results comparable to chain-of-thought (CoT) prompting method as shown in Feng et al. ().
While CoT outperforms the other methods, there is an increase in the serial computation due to its step-by-step reasoning approach. On the other hand, Pause Tokens and regularization primarily boost parallel computation, enhancing the model’s expressivity for these types of problems without significantly increasing serial computation.
The results presented in Figure 7 are based on models with a 5-layer configuration, optimized using the best-performing number of pause tokens. We find that the number of pause-tokens that works best for a particular model and layer config vary a bit. For example, the best working Seq-VCR + Pause model uses 7 pause tokens in case of the 5 operator task in Arithmetic Expression while 4 pause tokens works the best for the LIS task with 80 as input sequence length. Through Figure 8 we show how the different models perform by itself in the most complex task in each dataset. We also clearly see that more depth of the model helps every method in solving the task better. Another conclusion we can draw from figures 7 8 is that even though the combination of Seq-VCR with Pause Tokens works the best (almost close to CoT method), Seq-VCR gets similar test accuracy without the use of Pause Tokens in both datasets. All results in figure 78 are over 3 seed runs. Detailed ablation studies, exploring variations in the number of pause tokens and layers, can be found in the appendix.

## 5 Conclusion

Despite the remarkable success of large transformer based language models, they exhibit deficiencies in multi-step reasoning tasks. One proposed remedy is the application of Chain-of-Thought supervision, which, although effective, entails considerable cost. Conversely, there has been limited exploration into enhancing the capacity of the pre-trained model representation itself without explicit supervision. Our analysis indicates that intermediate layer representation collapse detrimentally affects the computation of intermediate information essential for solving arithmetic reasoning tasks. Our proposed regularization technique enhances the model’s representation capacity, thereby improving its ability to solve these tasks. We think that further research is required to enhance the capability of pre-training as well, where improved representation learning will enhance the model with greater capabilities.

#### Acknowledgments

We acknowledge the support of the Canada CIFAR AI Chair Program and IVADO. We thank Mila and Compute Canada for providing computational resources.

## Appendix A Appendix

### A.1 Experiments

**Table 2: Summary of Hyper-parameters, Compute Resources, and Time for Experiments. Each experiment was run on same configurations of GPUs, with the total time dependent on the number of layers and input parameters such as Dataset, task complexity.**
| Parameter | Value/Details |
| --- | --- |
| Learning Rate | 0.0001 |
| Batch Size | 128 |
| Optimizer | AdamW |
| Dropout | 0.1 |
| Attn. Heads | 4 |
| Epochs | 100 |
| Number of Layers | 4, 5, 6 |
| Input Sequence Length | 50, 80, 100 |
| # of Operators | 4, 5, 6 |
| Total Compute Resources | 1 32GB GPU, 6CPU, 32 GB RAM |
| Total Job Time | max 24 Hrs |

Figure: (a) Test accuracy on Arithmetic Expressions Dataset
Refer to caption: extracted/5970711/figures/arithmetic_accuracy_vs_layers.png

Both plots in Figure9 demonstrate the trade-offs between model complexity (in terms of layers) and the complexity of the input (number of operators or sequence length).
The model performs best when handling simpler configurations, such as fewer operators or shorter sequences, and tends to struggle with more complex setups as the number of operators or sequence length increases. Moreover, Increasing the number of layers seems to mitigate performance drops to some extent, especially in simpler cases, but it cannot fully offset the challenges posed by more complex inputs. All results in figure 9 are over 3 seed runs.

Figure: Figure 10: Ablation of Varying Pause Tokens and Comparing Vanilla Vs Seq-VCR + Pause Tokens
Refer to caption: extracted/5970711/figures/Vanilla_pause_tokens.png

Figure10 shows that Seq-VCR + Pause model consistently achieves high accuracy (around 0.8) across all numbers of pause tokens, indicating that the addition of pause tokens does not significantly affect its performance.
The Vanilla model, on the other hand, shows much lower accuracy across all settings, with a slight decrease in accuracy as the number of pause tokens increases (especially noticeable at 4, 6, and 8 pause tokens).

**Table 3: Dataset statistics. Size refers to the training set. The number of input, output, and intermediate chain of thought tokens are median values on the validation set. The number of tokens are based on the GPT-2 tokenizer, and a special ending symbol is counted for both intermediate tokens and output tokens.**
| Dataset | Task | Size | # Input<br>Tokens | # CoT<br>Tokens | # Output<br>Tokens |
| --- | --- | --- | --- | --- | --- |
| Multiplication | 4 x 4 Mult | 808k | 9 | 47 | 8 |
| 5 x 5 Mult | 808k | 11 | 75 | 10 |  |
| Arithmetic<br>Expression | 4 Operators | 1M | 19 | 24 | 1 |
| 5 Operators | 1M | 23 | 40 | 1 |  |
| 6 Operators | 1M | 27 | 60 | 1 |  |
| LIS | Seq Length 50 | 1M | 50 | 50 | 1 |
| Seq Length 80 | 1M | 80 | 80 | 1 |  |
| Seq Length 100 | 1M | 100 | 100 | 1 |  |

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 8
presentation: 3
questions: Same as the weaknesses section.
soundness: 3
strengths: - The paper presents a novel regularization technique that improves the model's performance in several reasoning tasks
- The paper presents detailed analysis of the experimental results, showcasing how exactly the regularization techniques affects the diversity of representations, the learning dynamics, as well as the digit-by-digit accuracy on multiplication tasks.
summary: This paper proposes a regularization technique for preventing representation collapse across the intermediate representations of a deep sequence model. Their results show that 1. the regularization technique increases matrix entropy (low matrix entropy = representation collapse) and 2. when pause tokens are added the language model significantly improved in performance for 4x4 and 5x5 arithmetic tasks.
weaknesses: - The effect of the regularization technique was only studied for a relatively narrow domain of tasks, and it would be interesting to understand its effect on more general language benchmarks as well.
- Slightly more contextualization on how exactly pause tokens are incorporated would assist readers in understanding the work more easily as it is also a core part of what is being proposed in this work.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: A discussion that clarified the improvement versus CoT would improve the significance, whether clearly establishing the speedup with Seq-VCR or showing its better generalization/scaling.
soundness: 3
strengths: The details of the method to seem to be heavily inspired by VICReg, but so far as I can judge, the application of it to the sequence/Transformer is original. The method is, in theory, computationally attractive compared to CoT and the results are fairly compelling. 

The paper is clearly written and the quality of the presentation is moderately high.
summary: This paper focuses on the performance of decoder-only models on tasks such as multi-digit mathematical reasoning that require a series of immediate representations. They hypothesize representation collapse of intermediate layers as a key contributor to this poor performance, preventing effective storage of the intermediate steps necessary for these kinds of tasks. While chain of thought reasoning can be effective in counteracting this collapse and performing well on such tasks, the proposed approach seeks to increase entropy among intermediate layers and achieve similar performance with at a reduced computational cost. Formulated in terms of alpha-order matrix-based entropy, the formulate a regularization term which aims are increasing variance and decreasing covariance in the intermediate representations. Additionally, pause tokens are included in the method. Results on three kinds of tests are presented – computing arithmetic expressions, identifying the longest increasing integer subsequence, and performing multiplication of 4 digit or 5 digit numbers. Performance with the regularization term and pause tokens leads to performance which approaches chain of thought on most tests, and regularization performs well on its own for the simpler tasks.
weaknesses: The result in Table 1 naturally provokes a question: This and several previous studies show that GPT-2 with CoT performs remarkably well, but this is actually more difficult to achieve in larger models. What is the evidence/argument that the Seq-VCR approach will scale better with model size than CoT? Figure 8 hints at this but it doesn’t clearly address it.

The speedup vs CoT is intuitively reasonable but it would have been nice to see performance numbers as in the cite Deng 2024 paper.

Similarly, it would be helpful to understand the amount of hyperparameter optimization necessary for, e.g., identifying the number of pause tokens used to obtain the best results. Do the number of pause tokens necessary correlate with, e.g., task complexity?

For completeness, it would be nice to see CoT in figures 7 and 8.

---

confidence: 5
contribution: 4
initial_rating: 5
presentation: 1
questions: (1) Pause tokens are a crucial part of the author's technique, but at no point do the authors describe where, and how, the pause tokens are added to the input.  

(2) Representational collapsed supposedly happens in the intermediate layers of the transformer, and yet the Lseq-VCR loss term is only applied to the final layer.  (Line 225).  Shouldn't it be applied to the intermediate layers, where you measure the entropy?  Why not?

(3) Equation (3) introduces $\lambda_1$ and $\lambda_2$ as hyperparameters, but the paper fails to say what they are set to. 

(4) What batch size is used for computing the covariance matrix?

(5) Equation 3 computes the covariance matrix only across the batch dimension.  Why?  In a transformer, you could potentially use the length dimension as well, which would drastically increase the effective batch size.  Did you do an ablation study which showed that to be ineffective?

(6) How is the projection layer $f_{proj}$ trained?

(7) For GPT2 on multiplication, you fine-tune a pre-trained GPT2 model, despite the fact that the pre-trained GPT2 has no real multiplication ability to start with.  Why bother with a pre-trained model, instead of just training from scratch, as you do with minGPT on arithmetic expressions?
soundness: 2
strengths: The application of VICReg to language models is novel as far as I know, and the experimental results are very compelling.  This could potentially be a high-impact paper in improving the reasoning capabilities of LLMs.
summary: Background: variance-covariance regularization (VICReg/VCReg) is a technique that was pioneered in vision models.  Given a batch of inputs, the technique uses an NN to encode those inputs to a batch of embedding vectors, and then computes a covariance matrix for the embedding vectors.  It introduces two losses based the covariance matrix: (a) the variance loss ensures that every dimension of the embedding vector has different values, distributed across the batch, and (b) the covariance loss ensures that different dimensions are not correlated.  In vision models, these two losses guard against representational collapse.  

The authors of this paper adapt VICReg from the vision domain to transformer-based language models.  They show that when combined with pause tokens, VICReg (now renamed to Seq-VCR) produces large improvements in several tasks that LLMs are usually very bad on -- multidigit arithmetic, arithmetic expressions, and a longest-increasing-subsequence task.
weaknesses: Unfortunately, the paper itself is hastily and sloppily written, and difficult to follow in places.  I had numerous questions when reading it that are not addressed by the text.  The current draft does not contain all of the information necessary to replicate these experiments, and does not discuss many of the crucial design decisions.  The authors claim that there is an "Appendix A", but neglected to provide any supplementary material.  See "questions" section below for specific questions.

One of the author's central claims is that transformers suffer from representational collapse, but I do not think that they adequately make that point based on the experimental evidence.  There are only two entropy charts in Figure 2, which cover only two narrow tasks.  On one of those charts (a) the collapse seems minimal at best, while on the other (b) the addition of pause tokens (the second key technique that the authors propose) actually increases collapse, rather than decreasing it.  I would need to see a much larger set of studies, over a variety of different tasks, including general language modeling tasks (translation etc.) to fully buy the author's argument about collapse.  If the authors did such a study, however, it would be a significant breakthrough.

Similarly, I would like to know what the effects of VICReg are on more general language modeling tasks.  If the technique helps the model multiply 5-digit numbers after fine-tuning, but otherwise degrades peformance on most other language modeling tasks, then the technique is useless.  Because the authors do not perform this ablation, it is impossible for me to evaluate whether this is a high-impact advance over SOTA, or a trivial result.

Finally, the use of pause tokens is interesting, but also seems haphazard.  They authors themselves admit that the number of pause tokens is task-specific.  To employ this technique more widely, I would need to see a more comprehensive test of where, how many, and under what circumstances pause tokens should be added.

More specific criticisms:

Equation (3) defines the Seq-VCR loss.  The text of the paper claims that it is "inspired by" prior work, and cites such work appropriately, but it is more than just "inspired".  Equation (3) is lifted almost verbatim from the orginal VICReg (Bardes 2021) and VCReg (Zhu 2023) papers, and the authors need to be crystal clear about the source of that equation.   

(As a minor nit, it is unclear to me whether or not the covariance term in equation (3) should have an additional 1/(d-1) factor; VICReg has the term, while VCReg does not.  I would have appreciated it if the authors explained why they chose one version over the other.)

For further clarity, the authors should also devote a few lines to defining how the covariance matrix C is computed; as is done in other papers.  Otherwise, it can easily be confused with the cross-correlation matrix of the Barlow twins technique, which the authors also cite as inspiration.

---

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: see the weaknesses section.
soundness: 2
strengths: 1. The identified representation collapse is quite interesting

2. The proposed method, including Seq-VCR regularization loss and pause tokens, demonstrates novelty and effectiveness based on the experimental results.
summary: This work identifies representation collapse in the LLM intermediate layers as a key factor limiting their arithmetic reasoning capabilities. The paper proposes sequential variance-covariance regularization (Seq-VCR). It then combines Seq-VCR with pause tokens to prevent the representation collapse. Experiments on GPT-2-small and minGPT demonstrate the effectiveness in improving accuracy on arithmetic reasoning.
weaknesses: 1. The representation collapse experiment was conducted only on GPT-2. I am curious whether this phenomenon occurs in more recent and larger LLMs, such as LLaMA 3 or LLaMA 3.1. The authors should either include additional experiments or provide a theoretical analysis to demonstrate that this is not an isolated case.

2. While the proposed Seq-VCR regularization loss has been shown to be effective in arithmetic reasoning tasks, I wonder whether adding this loss after the next token prediction loss would impact the LLM's performance on other tasks  (e.g., math reasoning and general MMLU). If it does have an effect, then this method may not be widely applicable. I encourage the authors to discuss this point.

Meta-review: This paper proposes a regularization technique for preventing representation collapse across the intermediate representations of a deep sequence model. Their results show that 1. the regularization technique increases matrix entropy (low matrix entropy = representation collapse) and 2. when pause tokens are added the language model significantly improved in performance for 4x4 and 5x5 arithmetic tasks.

The strengths of this paper are:
* simple technique 
* effectiveness on some simple reasoning tasks the authors experimented on

Weakness of this paper:
* unsure about how this technique would generate to leading LLMs (though the author added experiments to LLama during rebuttal)
* unsure about how this technique performs on more complex reasoning tasks.
* This method requires training, hence should be compared to any other reasoning augmentation method that also requires training.

Decision: Accept (Poster)