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

We would like to thank the reviewer for the time and valuable feedback. Below, we address the concerns raised.

1) Representation Collapse on larger LLMs: We have conducted additional large scale experiments and scaling analysis of collapse of different size of LLama models based on the reviewer's feedback. Please check the Generic Response (1) for more details.

2) Expanding Seq-VCR to broader language tasks: We appreciate the opportunity to address this valid concern. In response, we conducted additional experiments, which along with other details are discussed in the Generic Response. Please refer to Generic Response (2) for further details.

We hope we addressed all your concerns and please let us know if anything else is unclear. We would like to request the reviewer to consider increasing the score.

Could you provide the standard evaluation on general benchmarks, such as MMLU/GSM8K accuracy?

We would like to thank the reviewer for their feedback. We conducted additional experiments to fine-tune the GPT-2 Small model on an augmented version of the GSM8K dataset[1], both with and without Seq-VCR, Pause (2), and CoT. The results are shown below.

For this dataset, we observe slight performance improvement when applying pauses, Seq-VCR without Pause, and slightly more performance improvement when using Seq-VCR with pauses.

In addition to GSM8k, we have also included some preliminary pretraining results in Appendix L by training GPT2-Small on the C4 dataset. Seq-VCR models maintain performance on validation perplexity compared to the vanilla baseline model, while still increasing representation entropy (hence reducing collapse).

In summary, these two additional experiments show that our method:

• Does not penalize the performance on more generic language tasks

• Increases representation entropy (reduces representation collapse) in all cases

• Improves performance on mathematical reasoning tasks

[1] Deng, Yuntian, et al. "Implicit chain of thought reasoning via knowledge distillation." arXiv preprint arXiv:2311.01460 (2023).

Thank you very much for the detailed review and feedback. We address all your concerns below and will clarify in the paper accordingly. Please let us know if anything is still unclear.

-Appendix included in the main text:

We have included all supplementary materials directly in the main PDF document. We did not split it into two separate PDFs, as this was not explicitly required by the ICLR guidelines, and we do not have large zip files to share. The revised document contains several additional appendix sections.

-About representational collapse & the effects of Seq-VCR on more general language modelling tasks:

Please check the generic response section (1)

-About pause tokens (Q1):

Please check the generic response section (2)

-Regularization on intermediate layers (Q2):

In our initial experiments, we applied the regularization across all layers and observed better performance when it was applied only to the last layer. We think this is because, when applied to the last layer, the regularization loss gradients update all layers, whereas applying it to intermediate layers only updates those layers. While we found that applying it to the last layer works best, other layers could potentially yield similar results. We have added this clarification as a footnote in Section 3.4 of the manuscript, highlighted in blue-colored text.

-Hyperparameters (Q3, Q4):

For multiplication tasks, we set $\lambda_1 = 1.0$ and $\lambda_2 = 0.004$, while for other tasks, we use $\lambda_1 = 0.1$ and $\lambda_2 = 0.5$.

For computing the covariance matrix, we use a batch size of 32 for multiplication tasks and 128 for other tasks. These details have been added to Appendix A of the revised manuscript.

-Computation of the covariance matrix across both the batch and length dimensions (Q5):

New experiments: Thank you for pointing this out. We indeed missed this ablation when increasing the effective batch size. Based on your comment, we conducted an ablation study on the 5x5 digit multiplication task by fine-tuning GPT2-small. We compared the performance of calculating the covariance matrix across the batch dimension alone vs. across both the batch and length dimensions.

Results: We found that there are no significant differences in accuracy across multiple hyperparameter settings (ranging from 98% to 99% on the 5x5 digit multiplication task). However, the computation time is, on average, n times larger (depending on the sequence length, n) when calculating the covariance matrix across both the batch and length dimensions. Full results can be found in Appendix F.

-How projection layer $f_{proj}$, is trained (Q6):

The projection layer, $f_{proj}$, is added over the representation before calculating the loss over the projected representations. It is trained end-to-end exclusively using the Seq-VCR loss (Equation 3). We clarified this point in Section 3.4 of the paper in blue-colored text.

-Pre-trained models vs training from scratch (Q7):

New experiments: Based on your comment, we run more experiments to emphasize the advantage of training pre-trained models. We trained a gpt2 model from scratch on 5x5 digit multiplication tasks over multiple hyperparameter settings.

Results: Training models from scratch is more sensitive to hyperparameter choices such as batch size, hence your suggested way of increasing effective batch size worked better. For the full analysis, please check Table 6 in Appendix G.

-About equation 3:

The main contribution is to apply this regularization in Transformers on text-domain as the original idea was proposed for the image representation learning. We acknowledge that the equation is the same as VICReg[2] and therefore we use the same normalizing factors as in VICReg.

--Covariance Matrix Computation:

We computed the covariance matrix, C based on equation 3 in VICREG [2], so it's not the cross-correlation matrix as in the Barlow twins.

[1] Goyal, Sachin, et al. "Think before you speak: Training language models with pause tokens." arXiv preprint arXiv:2310.02226 (2023).

[2] Bardes, Adrien, Jean Ponce, and Yann LeCun. "Vicreg: Variance-invariance-covariance regularization for self-supervised learning." arXiv preprint arXiv:2105.04906 (2021).

We hope we have thoroughly addressed all your concerns and clarified any ambiguities. We kindly request the reviewer to consider revisiting the evaluation and potentially increasing the score based on these updates.

I would like to thank the authors for responding to my questions, running additional experiments, and for updating the paper. I especially appreciate the fact that modifications to the paper were done in blue, which makes it easy to see what was changed.

(1) Unfortunately, some of my original criticisms have not been addressed. The most glaring weakness of the paper, by far, is the fact that it does not measure the effect of SeqVCR on general-purpose language modeling. 5x5 digit multiplication problems are an extremely narrow downstream task, and a technique which improves performance on that task, but only at the expense of general language modeling, is useless. To clarify, I do not expect the authors to pre-train a large-scale LLM! It would be sufficient to train a small (e.g. 100M parameter) model from scratch on a general dataset like C4, and measure perplexity for next-token prediction. The purpose of the experiment would be merely to show that VCReg does not hurt performance on general LM tasks -- if the perplexity is the same, then your technique is a winner.

(2) A second weakness of the paper is that it conflates the effects of two completely different techniques -- SeqVCR and pause tokens. For the 5x5 digit multiplication task, the two seem to work together synergistically. However, in Appendix C. Figure 10, the best performing model seems to be the one with no pause tokens at all! And pause tokens seem to be particularly harmful to the vanilla model, for reasons that are unexplained. Given the weak results, I honestly feel that pause tokens are almost a distraction here.

(3) Thank you for running the additional experiments in Appendix G, which is actually a very interesting result. The authors write "training models from scratch is more sensitive to hyperparameter choices such as batch size". However, that's not the conclusion that I would draw from this experiment. By calculating the covariance matrix along sequence length, SeqVCR is ensuring that digits at different positions in the sequence have diverse representations. By using only the batch dimension, the SeqVCR loss improves diversity only among digits which occupy the same position. The fact that there is such a huge improvement in the first case says something interesting about the 5x5 multiplication problem.

(Conclusion) To be honest, I have mixed feelings about this paper. On the one hand, representation collapse is important, and SeqVCR in particular is an interesting technique. On the other hand, I don't think the current paper really does it justice. This has the potential to be a high-quality, high-impact paper if the authors properly explored the effect of SeqVCR on multiple tasks, including general language modeling and translation, did more ablations about the batch dimension vs. batch-dim + sequence-dim issue (Appendix G) on various tasks, etc.

Instead, the authors restrict their experiments to extremely narrow downstream tasks, and then try to present pause tokens as an alternative to CoT, when pause tokens seem to have minimal impact on anything other than multi-digit multiplication.

Additional comments:

Figure 6 -- Should have an additional accuracy bar for SeqVCR + Pause. I assume SeqVCR + Pause performs comparably to CoT because of Table 1, but that should also be shown in Figure 6.

Specific criticisms that have not yet been addressed (copied from before, with added emphasis).

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 in the text of the paper.

(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 (in the text of the paper) 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. To further clarify -- I know you are using the covariance matrix, and not the cross-correlation matrix. It might be helpful to other readers to have that spelled out by writing down the equation.

Thank you for your detailed review and positive feedback, we address your concerns below.

-Scalability of Seq-VCR compared to CoT: We appreciate this question. First, we would like to clarify that GPT3.5 and GPT4 results in Table 1 are from 5-shot prompted models, they are not fine-tuned models like GPT2. We believe that if GPT3.5 or GPT4 were fine-tuned with CoT they would achieve 100% results.

Regarding increased model sizes, Figure 8 shows that Seq-VCR is effective on larger (deeper) models. For additional results on scaling experiments, please check the generic response (1).

-Inclusion of CoT in Figure 8: Thank you for noticing that Figure 8 is missing the CoT numbers, we added them in the paper. They were 100% in all cases. Note, this is expected as CoT tokens carry more useful information than simple dummy <pause> tokens. However, <pause> combined with Seq-VCR serves as a competitive method while being computationally much cheaper and without using supervised CoT data for training.

-Speedup and Accuracy tradeoff for pause and CoT tokens: We carried out this analysis in response to the reviewer’s feedback. Our method is computationally much faster compared to the CoT method. For further details, please refer to the computational complexity analysis in Generic Response (2).

-Hyperparameter optimization: We conducted a manual search for hyperparameters, rather than performing exhaustive methods like grid search or random search. We found this approach sufficient to identify effective hyperparameters for our tasks. For the two coefficients, $\lambda_1$ and $\lambda_2$ in Eq. 3, we maintained their proportions close to those in [1]. Other hyperparameters, such as learning rate and batch size, are consistent with values used in related works [2, 3]. Further details are in Appendix A.

-# pause tokens vs. task complexity: Based on the reviewer’s recommendation, we conducted additional experiments on arithmetic tasks using a fixed 5-layer model, varying the number of pause tokens (2, 4, 6, 8) and task complexity. We did not observe a clear correlation between task complexity and the number of pause tokens. We hypothesize that this may be due to all pause tokens sharing the same embedding. Future work will investigate the impact of using different embeddings for each pause token. These details are now included in Appendix B of the revised manuscript.

[1] Bardes, Adrien, Jean Ponce, and Yann LeCun. "Vicreg: Variance-invariance-covariance regularization for self-supervised learning." arXiv preprint arXiv:2105.04906 (2021).

[2] Deng, Yuntian, et al. "Implicit chain of thought reasoning via knowledge distillation." arXiv preprint arXiv:2311.01460 (2023).

[3] Feng, Guhao, et al. "Towards revealing the mystery behind chain of thought: a theoretical perspective." Advances in Neural Information Processing Systems 36 (2024).

We hope we addressed all of your concerns and would be very happy to answer any followup questions. We would like to request the reviewer to consider increasing the score.

Dear Authors,

May I please ask you to clarify why the Vanilla model achieves a near perfect accuracy in Figure 4 (4 arithmetic operators) while slightly higher than 0.20 (zero pause tokens) in Figure 10 (again 4 arithmetic operators, the light blue colour)? Am I misreading it?

Dear Reviewer DQGy,

As we approach the end of the discussion period, we’d like to follow up on our previous responses to ensure they fully addressed your comments. If you have any additional questions or concerns, we’re happy to provide further clarification.

Thank you very much for your positive feedback on our work. We address your concerns in order below.

-The effect of the regularisation technique on the General Language Tasks: This is a valid point raised by the reviewer. Additional Experiments and discussion are in Generic Response (1). Please check that.

-Contextualizition of pause tokens: We really appreciate this feedback. We further discussed about pause token in the generic response (3) and also updating manuscript.

We are happy to address if the reviewer has any further queries.