On Exact Bit-level Reversible Transformers Without Changing Architecture

The authors thank all four reviewers for their appreciation of the novelty, simplicity, and effectiveness of our BDIA training technique for transformers. Notably, reviewer 7zwY states that the novelty and effectiveness of the approach make this a strong paper. Reviewer tmy4 states minimally changes the architecture allows the use of this method for finetuning, which is a creative and powerful use case. We would greatly appreciate the reviewers checking our detailed rebuttal responses and reconsidering their scores for our paper.

(1) Clarify the theoretical underpinnings of the reversibility claim explicitly.

In the revision, we will introduce a proposition and a proof in Subsection 4.3 to explicitly formulate the reversibility claim. Basically, in the proof, we will show that with quantization and the side information $s_{k-1}$, $x_{k-1}$ can be reconstructed exactly from $x_k$ and $x_{k+1}$ in online back-propagation, which only involves simple algebra. The side information $s_{k-1}$ essentially accounts for the binary quantization loss introduced by the configuration $\gamma_k\in \\{\pm 0.5\\}$.

(2) Include sensitivity analysis regarding quantization precision

As also suggested by reviewers 7ZwY and 8yrY, we fully agree that conducting the ablation study on the impact of different quantization levels would be really helpful to understand the behavior of BDIA. A complete sensitivity analysis would require additional investigation, but we have performed additional experiments for testing the performance of both ViT and BDIA-ViT for the quantization levels of $\\{5, 0, -2\\}$ over CIFAR10. In principle, the quantization level of $l=-2$ provides very coarse quantization effect and should lead to degraded performance.

As shown in the table below, when the quantization level goes from 5 to -2, the performance of ViT and BDIA-ViT decreases as expected. One can also observe from the table that when $l\in \\{5, 0\\}$, BDIA-ViT performs consistently better than ViT. For the case of $l=-2$, ViT performs slightly better than BDIA-ViT. This likely is because the very coarse quantization dominates the performance and negatively affects the BDIA-training technique.

(3) Regarding the relationship between quantization precision and reversibility.

In principle, the quantization precision is independent of the reversibility of BDIA-transformer. The reversibility of BDIA-transformer is guaranteed if the side information is stored for any particular quantization level. As we mentioned above, the quantization precision has a noticable impact on the validation performance. We will clarify the above point in the revision.

(4) Consider additional experiments explicitly varying model size and data to highlight potential trade-offs

Many thanks for the comments. We have conducted additional experiments by varying the number of transformer blocks in ViT and BDIA-ViT when trained over CIFAR10. As shown in the table below, BDIA-ViT performs consistently better than ViT for different numbers of transformer blocks. We can also observe that increasing the number of transformer blocks does not lead to better validation performance for a fixed dataset. Neural ``scaling" law suggest that the model size should scale along with the increasing size of the dataset in a log-log relationship.

We are not sure if we understand your suggestion on varying data for the additional experiments. It would be great if you could further elaborate on the above point so that we can perform appropriate additional experiments.

(5) Regarding the applicability of BDIA to other DNN blocks

Reviewer 8yrY suggested us to consider other transformer variants, such as linear transformer. Yes, we agree with your observation, although of course it would be necessary to empirically verify this. We have evaluated the BDIA technique for training vision-linformer over CIFAR10 by adopting the open source from https://github.com/pranavphoenix/VisionXformer/tree/main. Basically, we replaced the attention block in ViT by the linear attention block from the above open source in doing the experiment. The table below shows that BDIA again improves the validation accuracy considerably.

Due to very limited time for rebuttal, we don't have time to test other DNN blocks. It would be of great interest to further investigate the effectiveness of BDIA for other DNN blocks such as longformer, Mamba blocks, and convolutional networks in the future. We hope our work provides a motivation for the above research directions.

The authors thank all four reviewers for their appreciation of the novelty, simplicity, and effectiveness of our BDIA training technique for transformers. Notably, reviewer 7zwY states that the novelty and effectiveness of the approach make this a strong paper. Reviewer tmy4 states minimally changes the architecture allows the use of this method for finetuning, which is a creative and powerful use case. We would greatly appreciate the reviewers checking our detailed rebuttal responses and reconsidering their scores for our paper.

(1) Regarding the impact of different quantization levels on performance.

We have performed additional experiments to test the performance of both ViT and BDIA-ViT for the quantization levels of $\{5, 0, -2\}$ over CIFAR10. In principle, the quantization level $l=-2$ provides very coarse quantization effect and should lead to degraded performance.

As shown in the table below, when the quantization level goes from 5 to -2, the performance of ViT and BDIA-ViT decreases as expected. One can also observe from the table that when $l=5, 0$, BDIA-ViT performs consistently better than ViT. For the case of $l=-2$, ViT performs slightly better than BDIA-ViT. This could be because the very coarse quantization dominates the performance and negatively affects the BDIA-training technique.

(2) Regarding a discussion on the trade-off between memory savings and computational overhead

In the revision, we will add a table indicating the average running time per epoch for training ViT, BDIA-ViT with online backpropagation, and BDIA-ViT without online backpropagation over CIFAR10. We will then explain that BDIA-ViT with online backpropagation saves training memory and improves validation performance at the cost of slightly longer training time than ViT. We will also include additional figures in the revision showing the convergence curves in terms of wall-clock time.

(3) Regarding the impact of quantization level on memory savings

Firstly, we note that the reversibility of BDIA-transformer is guaranteed when the side information is stored for any quantization level in the forward pass of the model. The memory for the side information is fixed for a particular transformer model and is independent of the quantization level.

Secondly, in principle, when a coarse quantization is performed to the intermediate activation values across the transformer blocks, less memory is needed for the quantized activation values. Therefore, BDIA-transformer saves more memory when a coarser quantization is performed. However, the quantization cannot be too coarse as the performance of BDIA-transformer will be degraded, as shown in the above table.

(4) Can BDIA be applied to other transformer variants?

Many thanks for the suggestions. We have evaluated BDIA for training vision-linformer over CIFAR10 by adopting the open source from https://github.com/pranavphoenix/VisionXformer/tree/main. Basically, we replaced the attention block in ViT by the linear attention block from the above open source in doing the experiment. The table below shows that BDIA again improves the validation accuracy.

We believe that BDIA can be applied to other transformer variants. Due to limited time for rebuttal, we did not have time to test other attention variants. It would be of great interest to further investigate the effectiveness of BDIA for those variants in the future.

(5) What are the potential limitations of BDIA when scaling to large models?

These large models are rumored to require thousands of GPUs to train, and we will not have access to such resources, so our response is speculative. Our hypothesis is that a large model usually has many transformer blocks. For example, Google states that GPT-4 has 120 transformer blocks. In principle, the regularization impact of the BDIA training technique increases along with the increasing number of transformer blocks since the number of ODE solvers parameterized by $\\{\gamma_{k}\\}_{k=1}^{K-1}$ in BDIA increases exponentially with $K$. We hypothesize that the validation performance of the large model would be improved noticeably by using the BDIA training technique.

One thing that is not yet clear to us is if BDIA is compatible directly from an engineering viewpoint with the distributed data parallel (DDP) framework, which is widely used for training large transformer models across many GPUs. Additional engineering work might be required for solving the above potential issue.

(6) Regarding literature review on linear attention mechanisms and memory-efficient self-attention

We will update the introduction accordingly to reflect the works you mentioned.

The authors thank all four reviewers for their appreciation of the novelty, simplicity, and effectiveness of our BDIA training technique for transformers. Notably, reviewer 7zwY states that the novelty and effectiveness of the approach make this a strong paper. Reviewer tmy4 states minimally changes the architecture allows the use of this method for finetuning, which is a creative and powerful use case. We would greatly appreciate the reviewers checking our detailed rebuttal responses and reconsidering their scores for our paper.

(1) Regarding the computational overhead introduced by BDIA-transformer.

In the revision we will add a table indicating the average running time per epoch for training ViT, BDIA-ViT with online backpropagation, and BDIA-ViT without online backpropagation over CIFAR10. We will then explain that BDIA-ViT with online backpropagation saves training memory and improves validation performance at the cost of slightly longer training time than ViT. We will also include additional figures in the revision showing the convergence curves in terms of wall-clock time.

(2) Regarding a discussion of RevViT having better memory efficiency than BDIA-transformer

We will elaborate in the paragraph at lines 369-376 in the experimental section regarding the benefit of RevViT with respect to memory efficiency.

(3) Regarding the impact of different quantization levels

As also suggested by reviewers 8yrY and tmy4, we fully agree that conducting the ablation study on the impact of different quantization levels would be really helpful for understanding the behaviors of the BDIA training technique. We have performed additional experiments to test the performance of both ViT and BDIA-ViT for the quantization levels of $\{5, 0, -2\}$ over CIFAR10. In principle, a quantization level of $l=-2$ provides a very coarse quantization effect and should lead to degraded performance.

As shown in the table below, when the quantization level goes from 5 until -2, the performance of ViT and BDIA-ViT decreases as expected. For $l=-2$, there is a large performance drop due to the very coarse quantization effect. One can also observe from the table that when $l\in \\{5, 0\\}$, BDIA-ViT performs consistently and significantly better than ViT. For the extreme case of $l=-2$, ViT performs slightly better than BDIA-ViT. This likely is because the very coarse quantization dominates the performance and negatively affects the BDIA-training technique.

(4) Regarding integrating BDIA-transformers with other memory-saving techniques

There is one scenario where the BDIA training technique can be nicely combined with gradient checkpointing. We note that some transformer architectures may have special bottleneck blocks working as dimensionality transition (similar to the bottleneck blocks in ResNet). In the above case, a certain number of the transformer blocks of the same dimensionality are present between every two bottlenecks for efficient representation learning.

In the above scenario, one cannot directly perform online backpropagation from the very top transformer block to the bottom one due to dimensionality inconsistency. In this case, one can apply gradient checkpointing to handle the bottleneck blocks, and apply BDIA to do online backpropogation for those transformer blocks of the same dimensionality between every two bottlenecks.

(5) What specific challenges did you encounter when implementing BDIA for different transformer architectures?

Many thanks for the comment. We will release our source code once the paper is accepted so that other researchers can easily implement the BDIA training technique based on our source code.

From our experience, BDIA without online backpropagation can be easily implemented for different transformer architectures by following the update expressions in the paper. The challenging part is to realize online back-propagation for BDIA, which we implemented by modifying the open source code for RevViT. Researchers who only want to improve the validation performance of a transformer can simply ignore the online-backpropagation part of BDIA.

(6) Regarding comparisons against other memory-saving techniques beyond RevViT.

Many thanks for the comment. We plan to investigate the performance of the middle-point reversible transformer in the revision, of which the update expression is given by

$x_{k+1}= x_{k-1}+ 2 \\left[\\textrm{FFN}_k( {x}_k+\\textrm{Atten}_k({x}_k)) + \\textrm{Atten}_k({x}_k) \\right], \\quad k=1,\ldots, K-1.$

It is expected that BDIA-transformer will still perform better than the middle-point reversible transformer because the latter one does not introduce any regularization into the neural network.

The authors thank all four reviewers for their appreciation of the novelty, simplicity, and effectiveness of our BDIA training technique for transformers. Notably, reviewer 7zwY states that the novelty and effectiveness of the approach make this a strong paper. Reviewer tmy4 states minimally changes the architecture allows the use of this method for finetuning, which is a creative and powerful use case. We would greatly appreciate the reviewers checking our detailed rebuttal responses and reconsidering their scores for our paper.

(1) Regarding the paper readability

Based on comments from all four reviewers, we will make the following changes in the revision to improve readability:

(a) We will shorten the abstract to make it more concise.

(b) As suggested by reviewer tmy4, we will modify Subsection 4.3 by including a proposition as well as a proof for that proposition. The proposition will relate to the reversibility of BDIA with quantization and side information.

(c) The pseudocode for the BDIA training method will be included in the paper to enable the readers to easily understand the method.

(d) A few additional ablation studies will be included in the revision. In particular, we will study the impact of dropout rates, number of transformer blocks, and quantization levels on the validation performance of BDIA-ViT. We note that all the above ablation studies demonstrate that BDIA-ViT performs consistently better than ViT in terms of validation accuracy. We will also present the computational overhead of BDIA-ViT with online back-propagation in the revision.

(2) Regarding the reproducibility and pseudocodes of BDIA

In the revision, we will include the pseudocode for the BDIA training method. We will also release the source code after the paper is accepted so that other researchers will be able to adapt our source code to different transformer-based tasks.

(3) Although the paper includes thorough mathematical derivations, these seem to be more aligned with concepts from residual networks (ResNets) rather than focusing specifically on transformers.

It is true that the derivations take a unified viewpoint, i.e., that architectures with skip connections can be regarded from a numerical integration point of view. Components of the particular schemes (ResNet, NeuralODE, Transformer, etc.) are simply regarded as black-box functions in this view. This unified viewpoint makes BDIA applicable to not only transformers but also other schemes. The reason we focus on transformers in the paper is that transformers are commonly employed in models that (i) are very large and (ii) are the state-of-the-art in terms of performance. More-over, related methods have not been studied in the context of transformers. To reflect the well-motivated comment by the reviewer we will further strengthen the interpretation of our approach for transformers specifically in the final version of the paper.

(4) Regarding the small datasets in the experiments

Many thanks for the comment. Firstly, as commented by reviewer tmy4, the BDIA training technique could be a good candidate for fine-tuning transformer-based models because BDIA retains the standard architecture in the inference procedure. Fine-tuning a pre-trained model often works with a dataset of small or intermediate size. Our experiment for fine-tuning GPT2 medium using a small dataset in the paper shows that BDIA produces promising results.

Secondly, even though the considered datasets in our experiments are not very large, all three different training tasks (i.e., image classification, NLG, language translation) in our paper show that the BDIA training technique improves the validation performance of transformer consistently and considerably, demonstrating the potential applicability of the BDIA training technique in different applications.

(5) Regarding impact of dropout

Firstly, to clarify, the experimental results of Table 1 and 2 in the paper are obtained by setting the dropout rate to be 0.1. We will revise the paper accordingly to clarify the ambiguity.

Secondly, we conducted additional experiments by setting the dropout rate to be 0.0 and 0.2. The table below demonstrates that the dropout and BDIA can work together to improve the validation performance.

(6) Regarding figure 2

We will update figure 2 to demonstrate that the reconstruction error for BDIA with quantization and side information is zero. That is, lossless online back-propagation is guaranteed with BDIA thanks to the quantization and side information.

(7) Regarding figure 1

We will revise the caption of figure 1 to better explain how it is generated.