On Exact Bit-level Reversible Transformers Without Changing Architectures

## Abstract

Abstract Various reversible deep neural networks (DNN) models have been proposed to reduce memory consumption in the training process. However, almost all existing reversible DNNs either require special non-standard architectures or are constructed by modifying existing DNN architectures considerably to enable reversibility. In this work we present the BDIA-transformer, which is an exact bit-level reversible transformer 1 1 1 In fact, by following the same principle, we can also design exact bit-level reversible ResNets without changing its architecture for inference. Here, we focus on the transformer because of its popularity in the field of generative artificial intelligence. that uses an unchanged standard architecture for inference. The basic idea is to first treat each transformer block as the Euler integration approximation for solving an ordinary differential equation (ODE) and then incorporate the technique of bidirectional integration approximation (BDIA) (originally designed for diffusion inversion) into the neural architecture, together with activation quantization to make it exactly bit-level reversible. In the training process, we let a hyper-parameter γ 𝛾 \gamma italic_γ in BDIA-transformer randomly take one of the two values { 0.5 , − 0.5 } 0.5 0.5 \{0.5,-0.5\} { 0.5 , - 0.5 } per training sample per transformer block for averaging every two consecutive integration approximations. As a result, BDIA-transformer can be viewed as training an ensemble of ODE solvers parameterized by a set of binary random variables, which regularizes the model and results in improved validation accuracy. Lightweight side information per transformer block is required to be stored in the forward process to account for binary quantization loss to enable exact bit-level reversibility. In the inference procedure, the expectation 𝔼 ⁢ ( γ ) = 0 𝔼 𝛾 0 \mathbb{E}(\gamma)=0 blackboard_E ( italic_γ ) = 0 is taken to make the resulting architectures of BDIA-transformer identical to transformers up to activation quantization. Our experiments in both image classification and language translation show that BDIA-transformers outperform their conventional counterparts significantly in terms of validation performance due to the regularization effect of the set of γ 𝛾 \gamma italic_γ random variables while also requiring considerably less training memory.

## 1 Introduction

An active research trend in deep learning is to scale up the size of both the deep neural networks (DNNs) and the training data, with the aim of obtaining universal machine-learning models that are capable of accomplishing various tasks. Examples are large language models (LLMs) such as GPT-4 [] and Llama 2 [], which can, for example, have informative and friendly conversations with humans, solve mathematical problems, and produce high-quality source codes for programming tasks. A major bottleneck for training those large DNNs is that they often require large on-chip memory and large inter-chip communication bandwidth across many chips to accommodate both the DNN model and the intermediate activations of the input data-batch in back-propagation [], which is referred to as *memory wall* in the literature.

One promising technique to alleviate the issue of memory wall is to design and train reversible DNNs [; ; ; ]. By doing so, the intermediate activations in the forward pass do not have to be stored in the memory to allow for back-propagation. Instead, they can be recomputed on-the-fly in the backward pass by exploiting the reversibility of the DNN, thus saving memory consumption by a large margin for deep DNNs. The procedure essentially reduces memory consumption at the cost of a reasonable amount of additional computation. The reduction in memory also has the potential to improve throughput by increasing the batch size.

Early reversible DNNs, such as NICE [], the methods of [], Real NVP [], and Inverse Autoregressive Flow [] are differentiable transformations that were aimed at facilitating generative modeling. For an overview of such early normalizing flows see []. Inspired by these works, a range of reversible residual models were proposed subsequently, including RevNet [],
Glow [], i-RevNet [], i-ResNet [], layerwise inversion [], Fourier-transformation based CNN inversion [], and Mintnet []. Another line of research enforces reversibility in deep learning from the perspective of ordinary differential equations (ODEs), which includes FFJORD [], leapfrog networks [], neural ODEs [], momentum residual networks [], and neural ODE inversion [].

Recently, the research on reversibility has moved to other types of neural networks. The authors in [; ] proposed reversible vision transformers (referred to as RevViT) due to the popularity of LLMs. [] utilized a reversible diffusion sampling method to optimize the noisy representation for the task of diffusion based image editing. To our best knowledge, all the above existing reversible DNNs either require non-standard architectures or are constructed by modifying the original DNN architectures considerably to enable reversibility.

In this paper, we propose BDIA-transformer, a new type of reversible transformer that uses an unchanged, standard architecture for the inference procedure. It is based on the bidirectional integration approximation (BDIA), which was recently proposed in [] to enable diffusion inversion for round-trip image editing. To be able to incorporate BDIA into transformers for online back-propagation, we follow the common practice of treating each transformer block as an Euler integration approximation for solving an ordinary differential equation (ODE).

We make two main contributions in this work. Firstly, we propose BDIA-transformers by introducing a random hyper-parameter $\gamma\in\{-0.5,0.5\}$ per transformer block per training sample to regularize the DNN models for improvement of validation performance. Each $\gamma$ parameter intends to average every two consecutive integration approximations. The training procedure becomes training of an ensemble of ODE solvers parameterized by a set of binary random variables. In the inference procedure, the expectation $\mathbb{E}[\gamma]=0$ is utilized, which reduces BDIA-transformers to conventional transformers. Therefore, our method for enforcing reversibility is more flexible than existing ones.

Secondly, we perform activation quantization to allow for exact bit-level reversibility of BDIA-transformers. Note that the special setup of the $\gamma$ values in the subset $\{-0.5,0.5\}$ when performing activation quantization leads to a 1 bit information loss per activation value per transformer block. Therefore, lightweight side information per transformer block needs to be stored during training to recover this 1 bit information loss. Despite this, the overall memory use is significantly reduced.

Experimental results on ViT for image classification and transformer for language translation show that
the BDIA technique significantly improves the validation performance over that of the corresponding baseline transformers and simultaneously reduces training memory significantly. The improved performance results from the model regularization imposed by a set of $\gamma$ random variables. Experiments on text prediction demonstrate that BDIA-GPT2 significantly alleviates the over-fitting issue of GPT2 when intentionally trained on a very small dataset. Our empirical study also indicates that RevViT from [] produces either inferior or comparable validation performance to that of its original counterparts.

## 2 Related Works

In recent years, various quantization strategies [; ; ] have been proposed in the training and/or inference processes of DNN models on low-precision devices. For instance, the work [] successfully performed quantization on DNN weights, activations, gradients and errors in the training and inference processes and obtained promising results. The recent work [] demonstrated that LLMs with a quantization of 1.58 bits per model parameter exhibit comparable performance to non-quantized models. In summary, it was found that the validation performance of DNN models that incorporate those quantization operations is comparable with that of conventional DNN models. In our work, we only need to apply activation quantization to enable exact bit-level reversibility in training BDIA-transformers.

We note that in general, model quantization and design of reversible DNN models are two complimentary strategies for reducing memory consumption in the training process, where the first strategy operates on the model parameters and the second one operates on the activation values. With the development of new model quantization methods such as [], advancing the research frontier of reversible DNN models has become of great interest.

## 3 Preliminary

Neural networks as ODEs:
[] highlighted the interpretation that
passing a hidden state across layers that add a correction to that state can be viewed as Euler integration. While that paper identified architectures including residual nets and normalizing flows as following this pattern, it is equally true of diffusion models and transformers. This interpretation emphasizes the importance of considering more accurate integration schemes []. The need for improved integration is particularly apparent in round-trip image editing in diffusion models, where significant integration error will result in unintended visible alterations to the image.

Diffusion sampling via solving ODE: Recently, the work [] proposed the BDIA technique to enable diffusion inversion for effective round-trip image editing. From a high level point of view, BDIA can be viewed as a time-reversible ODE solver.
Given an initial diffusion state $\boldsymbol{z}_{T}$ at time step $T$, the
diffusion-based sampling process for generating realisic images $\boldsymbol{z}_{\epsilon}$ at time $t=\epsilon>0$ can be realized by solving a probability ordinary differential equation (ODE)

$$ $\displaystyle d\boldsymbol{z}=\boldsymbol{d}(\boldsymbol{z},t)dt$ (1) $$

over the time interval $t\in[T,\epsilon]$. The gradient vector $\boldsymbol{d}(\boldsymbol{z},t)$ includes the output of a pre-trained DNN model with $(\boldsymbol{z}_{t},t)$ as its input. The common practice for solving the above ODE is to first discretize the continuous time interval $[T,\epsilon]$ properly into a set of timesteps $\{t_{i}|i=0,\ldots,N\}$ with $t_{0}=T$ and $t_{N}=\epsilon$, and then perform certain integration approximation per small time-interval sequentially to compute the final diffusion state $\boldsymbol{z}_{N}=\boldsymbol{z}_{\epsilon}$.

BDIA: Suppose we would like to estimate the next diffusion state $\boldsymbol{z}_{i+1}=\boldsymbol{z}_{t_{i+1}}$ by solving (1)
based on the recent information $(\boldsymbol{z}_{i},t_{i})$ and $(\boldsymbol{z}_{i-1},t_{i-1})$, where $\boldsymbol{z}_{j}=\boldsymbol{z}_{t_{j}}$ for $j=i-1,i$. The basic idea of BDIA is to compute $\boldsymbol{z}_{i+1}$ by performing both the forward integration approximation $\Delta(t_{i}\rightarrow t_{i+1}|\boldsymbol{z}_{i})$ $\left(\approx\int_{t_{i}}^{t_{i+1}}\boldsymbol{d}(\boldsymbol{z},t)dt\right)$ and the backward integration approximation $\Delta(t_{i}\rightarrow t_{i-1}|\boldsymbol{z}_{i})$ $\left(\approx-\int_{t_{i-1}}^{t_{i}}\boldsymbol{d}(\boldsymbol{z},t)dt\right)$ conditioned on $\boldsymbol{z}_{i}$. One popular method for implementing $\Delta(t_{i}\rightarrow t_{i+1}|\boldsymbol{z}_{i})$ and $\Delta(t_{i}\rightarrow t_{i-1}|\boldsymbol{z}_{i})$ in the literature of diffusion models is by employing the DDIM update expression (see [; ; ] for details). With the above two integration approximations, $\boldsymbol{z}_{i+1}$ can be expressed as

$$ $\displaystyle\boldsymbol{z}_{i+1}=$ $\displaystyle\boldsymbol{z}_{i-1}\underbrace{-(1-{\color[rgb]{0,0,1}\gamma})( \boldsymbol{z}_{i-1}-\boldsymbol{z}_{i})-{\color[rgb]{0,0,1}\gamma}\Delta(t_{i }\rightarrow t_{i-1}|\boldsymbol{z}_{i})}_{\approx\int_{t_{i-1}}^{t_{i}} \boldsymbol{d}(\boldsymbol{z},t)dt}+\underbrace{\Delta(t_{i}\rightarrow t_{i+1 }|\boldsymbol{z}_{i})}_{\approx\int_{t_{i}}^{t_{i+1}}\boldsymbol{d}( \boldsymbol{z},t)dt}$ (2) $\displaystyle=$ $\displaystyle\gamma\boldsymbol{z}_{i-1}+(1-\gamma)\boldsymbol{z}_{i}-\gamma \Delta(t_{i}\rightarrow t_{i-1}|\boldsymbol{z}_{i})+\Delta(t_{i}\rightarrow t_ {i+1}|\boldsymbol{z}_{i}),$ (3) $$

where $\gamma=(0,1]$ averages $\Delta(t_{i}\rightarrow t_{i-1}|\boldsymbol{z}_{i})$ and $(\boldsymbol{z}_{i-1}-\boldsymbol{z}_{i})$ for the time-slot $[t_{i-1},t_{i}]$. The minus sign in front of the two quantities are due to the reverse integration direction. The quantity $-(\boldsymbol{z}_{i-1}-\boldsymbol{z}_{i})$
is the previously computed integration approximation for $\int_{t_{i-1}}^{t_{i}}\boldsymbol{d}(\boldsymbol{z},t)dt$.
We note that negative $\gamma$ values are not applicable to the diffusion models considered in [].

On reversibility of BDIA: The update expression (3) is carefully designed in [] to enable diffusion inversion for round-trip image editing. By reformulating (3), $\boldsymbol{z}_{i-1}$ can be easily computed in terms of $(\boldsymbol{z}_{i},\boldsymbol{z}_{i+1})$.
We note that due to the nature of the floating-point datatype, there might be error accumulation in round-trip image reconstruction, where the corresponding diffusion states in the forward and reverse process are not identical. In practice, error accumulation of BDIA in diffusion inversion is not a big issue [] due to the fact that the number of timesteps is generally set to be small (e.g., 50 timesteps in either forward or reverse process) to make the time complexity reasonable.

One main difference between diffusion inversion and reversible transformers is that no gradient needs to be back-propagated in diffusion inversion for updating the DNN model. As a result, even if there is error accumulation in diffusion inversion, it is less severe than in reversible transformers where error accumulation in online back-propagation would slow down the training convergence or even make the training fail especially for very deep models like LLMs. In next section, we will explain how to design exact bit-level reversible transformers in the training process to avoid any error-accumulation while at the same time, maintaining the architectures of the transformer in the inference procedure.

## 4 Exact Bit-Level Reversible Transformers via BDIA

In this section, we first briefly review the transformer update expressions from the ODE viewpoint. We then propose the BDIA-transformer that enables exact bit-level reversibility with activation quantization. Specially, we will demonstrate how each of the two $\gamma$ values $\{-0.5,-0.5\}$ averages consecutive integration approximations. The training of a BDIA-transformer can then be interpreted as employing different ODE solvers for different training samples in a random manner. In addition, we explain why additional lightweight side-information is required to be stored to account for the binary quantization loss in online back-propagation.

### 4.1 Revisiting transformer update expression

A typical transformer block consists of the attention function (denoted as $\boldsymbol{f}(\cdot)$) and the function of feed-forward network (FFN) (denoted as $\boldsymbol{g}(\cdot)$), of which the trainable parameters are generally different for different block indices. Accordingly, the output $\boldsymbol{x}_{k+1}$ of the $k$th transformer block can be mathematically represented in terms of the input $\boldsymbol{x}_{k}$ as

$$ $\displaystyle\boldsymbol{x}_{k+1}$ $\displaystyle=\boldsymbol{x}_{k}+\underbrace{\boldsymbol{f}_{k}(\boldsymbol{x} _{k})+\boldsymbol{g}_{k}(\boldsymbol{x}_{k}+\boldsymbol{f}_{k}(\boldsymbol{x}_ {k}))}_{\boldsymbol{h}_{k}(\boldsymbol{x}_{k})},$ (4) $$

where we use $\boldsymbol{h}_{k}(\boldsymbol{x}_{k})$ to denote the overall residual quantity that includes both the attention and FFN functions. For simplicity, we omit the pre-normalisation operations in (4), since these in fact do not affect the design of BDIA-transformers later on.

It is well known from the literature [] that the $k$th forward step in (4) can be viewed as the Euler integration approximation of an ODE at timestep $t_{k}$:

$$ $\displaystyle\boldsymbol{x}_{k+1}=\boldsymbol{x}_{k}+\boldsymbol{h}_{k}( \boldsymbol{x}_{k})$ $\displaystyle=\boldsymbol{x}_{k}+\overbrace{\tilde{\boldsymbol{d}}(\boldsymbol {x}_{k},t_{k})(t_{k+1}-t_{k})}^{\Delta(t_{k}\rightarrow t_{k+1}|\boldsymbol{x} _{k})}$ (5) $\displaystyle\approx\boldsymbol{x}_{k}+\int_{t_{k}}^{t_{k+1}}\tilde{ \boldsymbol{d}}(\boldsymbol{x},t)dt,$ $$

where $\tilde{\boldsymbol{d}}(\boldsymbol{x}_{k},t_{k})$ denotes the gradient vector with $(\boldsymbol{x}_{k},t_{k})$ as the input. Both $\tilde{\boldsymbol{d}}(\boldsymbol{x}_{k},t_{k})$ and $(t_{k+1}-t_{k})$ are implicitly learned via the composite function $\boldsymbol{h}_{k}(\boldsymbol{x}_{k})$, which is alternatively denoted as $\Delta(t_{k}\rightarrow t_{k+1}|\boldsymbol{x}_{k})$.

As explained later on, we will introduce BDIA into the update expression of (5). Note that (5) is a general update expression that not only includes the transformer but also ResNet [] and its variants. In fact, the following analysis can be easily adapted to enable reversibility of ResNet variants.

### 4.2 BDIA-transformer without quantization

In this subsection, we first derive the update expressions of BDIA-transformer without quantization as an extension of (5), and then study the impact of the $\gamma$ values in $\{0.5,-0.5\}$. When $k=0$, $\boldsymbol{x}_{1}$ can be computed by following (5) as

$$ $\displaystyle\boldsymbol{x}_{1}$ $\displaystyle=\boldsymbol{x}_{0}+\boldsymbol{h}_{0}(\boldsymbol{x}_{0})= \boldsymbol{x}_{0}+\Delta(t_{0}\rightarrow t_{1}|\boldsymbol{x}_{0}).$ (6) $$

The update expression for $\boldsymbol{x}_{k+1}$ in the training process, $K-1\geq k\geq 1$, can be obtained by utilizing (2)-(3). Based on (5), we let

$$ $\displaystyle\Delta(t_{k}\rightarrow t_{k-1}|\boldsymbol{x}_{k})$ $\displaystyle=-\boldsymbol{h}_{k}(\boldsymbol{x}_{k})$ (7) $\displaystyle\Delta(t_{k}\rightarrow t_{k+1}|\boldsymbol{x}_{k})$ $\displaystyle=\boldsymbol{h}_{k}(\boldsymbol{x}_{k}).$ (8) $$

By combining (7)-(8) and (2)-(3) with $i=k$, the update expression $\boldsymbol{x}_{k+1}$, $K-1\geq k\geq 1$, can be represented as

$$ $\displaystyle\boldsymbol{x}_{k+1}$ $\displaystyle=\boldsymbol{x}_{k-1}+(1-\gamma)(\boldsymbol{x}_{k}-\boldsymbol{x }_{k-1})+(1+\gamma)\boldsymbol{h}_{k}(\boldsymbol{x}_{k})$ (9) $\displaystyle=\gamma\boldsymbol{x}_{k-1}+(1-\gamma)\boldsymbol{x}_{k}+(1+ \gamma)\boldsymbol{h}_{k}(\boldsymbol{x}_{k}),$ (10) $$

where $\gamma$ is recommended to take values from $\{0.5,-0.5\}$ with equal probability per training sample per transform block, the impact of which will be explained in detail in the following. This is different from the work of BDIA-based diffusion inversion [] where $\gamma$ has to be positive.

In the inference stage, $\mathbb{E}(\gamma)=0$ is taken to replace $\gamma$ in (10), which leads to a simpler update expression that only involves $(\boldsymbol{x}_{k},\boldsymbol{x}_{k+1})$:

$$ $\displaystyle\boldsymbol{x}_{k+1}=\boldsymbol{x}_{k}+\boldsymbol{h}_{k}( \boldsymbol{x}_{k}),$ (11) $$

which is identical to the original transformer update expression (4)-(5).

Figure: Figure 1: Validation performance of different ODE solvers parameterized by a single $\gamma$ parameter after training ViT and BDIA-ViT over CIFAR10. See Subsection 5.1 on how ViT and BDIA-ViT were trained. Each ODE solver in the inference procedure is realized by selecting $\gamma$ from $[-0.5,0.5]$, which is fixed across all the transformer blocks for the same input image. The validation performance of BDIA-ViT is more robust than that of ViT.
Refer to caption: x1.png

Impact of $\gamma$ parameter: We now study the impact of the $\gamma$ parameter in (10). When $\gamma=0.5$, it follows from (2)-(3) and (7)-(8) that the two integrations $\int_{t_{k-1}}^{t_{k}}\tilde{\boldsymbol{d}}(\boldsymbol{x}_{\tau},\tau)d\tau$ and $\int_{t_{k}}^{t_{k+1}}\tilde{\boldsymbol{d}}(\boldsymbol{x}_{\tau},\tau)d\tau$ of the transformer ODE are approximated as

$$ $\displaystyle\int_{t_{k-1}}^{t_{k}}\tilde{\boldsymbol{d}}(\boldsymbol{x}_{\tau },\tau)d\tau\approx 0.5(\boldsymbol{x}_{k}-\boldsymbol{x}_{k-1})+0.5 \boldsymbol{h}_{k}(\boldsymbol{x}_{k})$ (12) $\displaystyle\int_{t_{k}}^{t_{k+1}}\tilde{\boldsymbol{d}}(\boldsymbol{x}_{\tau },\tau)d\tau\approx\boldsymbol{h}_{k}(\boldsymbol{x}_{k}),$ (13) $$

where the integration over $[t_{k-1},t_{k}]$ is computed as the weighted average of two consecutive integration approximations: $(\boldsymbol{x}_{k}-\boldsymbol{x}_{k-1})$ and $\boldsymbol{h}_{k}(\boldsymbol{x}_{k})$.

When $\gamma=-0.5$, the two integrations $\int_{t_{k-1}}^{t_{k}}\tilde{\boldsymbol{d}}(\boldsymbol{x}_{\tau},\tau)d\tau$ and $\int_{t_{k}}^{t_{k+1}}\tilde{\boldsymbol{d}}(\boldsymbol{x}_{\tau},\tau)d\tau$ are approximated differently, given by

$$ $\displaystyle\int_{t_{k-1}}^{t_{k}}\tilde{\boldsymbol{d}}(\boldsymbol{x}_{\tau },\tau)d\tau\approx(\boldsymbol{x}_{k}-\boldsymbol{x}_{k-1})$ (14) $\displaystyle\int_{t_{k}}^{t_{k+1}}\tilde{\boldsymbol{d}}(\boldsymbol{x}_{\tau },\tau)d\tau\approx 0.5\boldsymbol{h}_{k}(\boldsymbol{x}_{k})+0.5(\boldsymbol{ x}_{k}-\boldsymbol{x}_{k-1}).$ (15) $$

In this case, the integration over $[t_{k},t_{k+1}]$ is computed by averaging $(\boldsymbol{x}_{k}-\boldsymbol{x}_{k-1})$ and $\boldsymbol{h}_{k}(\boldsymbol{x}_{k})$.

Training via an ensemble of ODE solvers: We use $\gamma_{k}$ to denote the random variable for the $k$th transformer block taking values in $\{0.5,-0.5\}$ with equal probability. The above analysis of (12)-(15) implies that each training sample goes through a particular ODE solver determined by sampling a set of $K-1$ random variables $\{\gamma_{k}\}_{k=1}^{K-1}$, which specifies the integration path from the bottom transformer block until the top one. Due to randomness, different training samples will go through different ODE solvers. There are in total $2^{K-1}$ different ODE solvers, where each one corresponds to a unique integration path across the $K$ transformer blocks. Intuitively speaking, if we assume that each of the individual ODE solvers is well behaved (i.e., its training loss decays to a small value), then a convex combination will also converge to a small value.

In principle, different ODE solvers can also be applied in the inference procedure. For instance, one can set $\gamma$ to be constant within $[-0.5,0.5]$ in (10) across all the transformer blocks for the same input. Fig. 1 demonstrates the validation performance of those different ODE solvers in the inference procedure for both trained BDIA-ViT and ViT over CIFAR10. It is clear that the validation performance of BDIA-ViT is much more insensitive to the single $\gamma$ parameter than that of ViT. The reason for the sensitivity of ViT to $\gamma$ in Fig. 1 is because only a single ODE solver is trained for ViT.

On similarity to Dropout technique: From a high level point of view, the training of BDIA-transformer is similar to the conventional dropout technique [] to a certain extent. Dropout essentially attempts to train an ensemble of subnetworks of the entire DNN model when minimizing the objective function while BDIA-transformer intends to train an ensemble of different ODE solvers. In the inference procedure, an average of the ensemble of subnetworks in Dropout is utilized while in BDIA-transformer, the expectation $\mathbb{E}[\gamma_{k}]=0$ is employed to replace $\gamma_{k}$.

###### Remark 1 .

If reversibility is not of concern, one can freely specify the values for the random variables $\{\gamma_{k}\}_{k=1}^{K-1}$ in the training process for performance improvement as long as its distribution is symmetric around 0 0 such that $\mathbb{E}[\gamma_{k}]=0$ for maintaining the original DNN architectures in the inference procedure. See Table 2 for ablation study of the impact of the $\{\gamma_{k}\}_{k=1}^{K-1}$ parameter on the validation performance when training BDIA-transformer for image classification.

###### Remark 1 .

### 4.3 On exact bit-level reversibility of BDIA-transformer with quantization

As explained in Section 1, one strategy for reducing memory consumption in training transformer models like LLMs is to perform online back-propagation. That is, the intermediate activation outputs from the top transformer block until the bottom one are computed online when performing back-propagation to update the model. In this subsection, we first discuss the reversibility issue of the update expression (10. We then consider performing activation quantization to enable exact bit-level reversibility. We note that lightweight side information is required to be stored per transformer block for lossless online back-propagation.

Figure: Figure 2: Demonstration of the accumulated reconstruction error by following (16) with the setup $\gamma_{k}\in\{0.5,-0.5\}$, $k=1,\ldots,K{-}1$, when training BDIA-GPT2 with 12 transformer blocks.
Refer to caption: x2.png

Limitation of the reversibility of (10): The update expression (10 is only theoretically reversible. That is, $\boldsymbol{x}_{k-1}$ can be computed in terms of $(\boldsymbol{x}_{k},\boldsymbol{x}_{k+1})$ as

$$ $\displaystyle\boldsymbol{x}_{k-1}=\boldsymbol{x}_{k+1}/\gamma_{k}-(1-\gamma_{k })/\gamma_{k}\boldsymbol{x}_{k}-(1+\gamma_{k})/\gamma_{k}\boldsymbol{h}_{k}( \boldsymbol{x}_{k}),$ (16) $$

where we use $\gamma_{k}$ to indicate that different transformer blocks have their respective random variables.
In practice, the setup $\gamma_{k}\in\{0.5,-0.5\}$, $k=1,\ldots,K-1$, would lead to non-negligible error accumulation especially for very deep transformer models. The factor $\frac{1}{\gamma_{k}}=\pm 2$ in front of $\boldsymbol{x}_{k+1}$ would amplify the error when $k$ decreases from $K-1$ to 1, making the online back-propagation unstable. Fig. 2 illustrates that the reconstruction error indeed increases significantly when applying the online-back-propagation from the top transformer block until the bottom one.

BDIA-transformer with quantization: To allow for lossless online back-propagation, we propose to perform activation quantization. In particular, we use $\mathcal{Q}_{l}[\cdot]$ to denote the quantization operation to the bit-level precision of $2^{-l}$, given by

$$ $\displaystyle\mathcal{Q}_{l}[y]=\textrm{round}[y/2^{-l}]2^{-l}.$ (17) $$

Upon introducing $\mathcal{Q}_{l}[\cdot]$, the new update expression for BDIA-transformer can be represented as

$$ $\displaystyle\boldsymbol{x}_{0}$ $\displaystyle\leftarrow\mathcal{Q}_{l}[\boldsymbol{x}_{0}],$ (18) $\displaystyle\boldsymbol{x}_{1}$ $\displaystyle=\boldsymbol{x}_{0}+\mathcal{Q}_{l}[\boldsymbol{h}_{0}( \boldsymbol{x}_{0})]$ (19) $\displaystyle\boldsymbol{s}_{k-1}[m]$ $\displaystyle=\left\{\begin{array}[]{ll}1&\textrm{if }\textrm{mod}(\boldsymbol {x}_{k-1}[m]/2^{-l},2)=1\\ 0&\textrm{otherwise}\end{array}\right.\quad k\geq 1$ (22) $\displaystyle\boldsymbol{x}_{k+1}$ $\displaystyle=\mathcal{Q}_{l}[\gamma_{k}(\boldsymbol{x}_{k-1}+\boldsymbol{s}_{ k-1}2^{-l})]+\mathcal{Q}_{l}[(1-\gamma_{k})\boldsymbol{x}_{k}+(1+\gamma_{k}) \boldsymbol{h}_{k}(\boldsymbol{x}_{k})]\quad k\geq 1,$ (23) $$

where $\gamma_{k}\in\{0.5,-0.5\}$, and $\boldsymbol{x}_{k-1}[m]$ denotes the $m$th element of $\boldsymbol{x}_{k-1}$. The $m$th element $\boldsymbol{s}_{k-1}[m]$ indicates if the integer value $\boldsymbol{x}_{k-1}[m]/2^{-l}$ is odd or not. It is immediate from (18)-(23) that
$\boldsymbol{x}_{k}=\mathcal{Q}_{l}[\boldsymbol{x}_{k}]$ for all $K\geq k\geq 0$. That is, all the intermediate activation outputs $\{\boldsymbol{x}_{k}\}_{k=0}^{K}$ have fix-point precision of $2^{-l}$.

Again in the inference procedure, we replace $\gamma_{k}$ in (23) by $\mathbb{E}(\gamma_{k})=0$. As a result, the update expression (23) can be simplified to be

$$ $\displaystyle\boldsymbol{x}_{k+1}$ $\displaystyle=\mathcal{Q}_{l}[\boldsymbol{x}_{k}+\boldsymbol{h}_{k}( \boldsymbol{x}_{k})]\quad k\geq 1.$ (24) $$

The only difference of (24) w.r.t. the original transformer update expression (4) is that the quantization operation $\mathcal{Q}_{l}[\cdot]$ is performed for each activation output.

On reversibility of (23) by storing lightweight side information: We now consider recovering $\boldsymbol{x}_{k-1}$ from $(\boldsymbol{x}_{k},\boldsymbol{x}_{k+1})$ by utilizing (22)-(23). By using the fact that $\gamma_{k}\in\{0.5,-0.5\}$ and the definition for $\boldsymbol{s}_{k-1}$, we can easily conclude that the quantity $\mathcal{Q}_{l}[\gamma_{k}(\boldsymbol{x}_{k-1}+\boldsymbol{s}_{k-1}2^{-l})]$ in (23) can be alternatively represented as

$$ $\displaystyle\mathcal{Q}_{l}[\gamma_{k}(\boldsymbol{x}_{k-1}+\boldsymbol{s}_{k -1}2^{-l})]=\gamma_{k}(\boldsymbol{x}_{k-1}+\boldsymbol{s}_{k-1}2^{-l}).$ (25) $$

That is, the quantization operation has no effect on $\gamma_{k}(\boldsymbol{x}_{k-1}+\boldsymbol{s}_{k-1}2^{-l})$. This is because the vector $\boldsymbol{s}_{k-1}$ essentially captures the 1-bit quantization loss of $\mathcal{Q}_{l}[\gamma_{k}\boldsymbol{x}_{k-1}]$ per element.

Suppose in each forward pass in the training process, all the side information $\{\boldsymbol{s}_{k-1}\}_{k=1}^{k=K-1}$ are stored in the memory. When we perform online back-propagation, each $\boldsymbol{x}_{k-1}$ can be reconstructed losslessly in the form of

$$ $\displaystyle\boldsymbol{x}_{k-1}$ $\displaystyle=\frac{1}{\gamma_{k}}\boldsymbol{x}_{k+1}-\boldsymbol{s}_{k-1}2^{ -l}-\frac{1}{\gamma_{k}}\mathcal{Q}_{l}[(1-\gamma_{k})\boldsymbol{x}_{k}+(1+ \gamma_{k})\boldsymbol{h}_{k}(\boldsymbol{x}_{k})],\;\;\;\;k\geq 1.$ (26) $$

Consequently, the computed gradient in the online back-propagation will not be drifted from the ground truth, which is desirable in very deep transformer models.

###### Remark 2 .

In principle, one can also make BDIA-transformer with $\{\gamma_{k}\in\pm 0.25\}_{k=1}^{K-1}$ to be exactly bit-level reversible. In this scenario, side information of 2 bits per activation value per transformer block is required to be stored to account for the quantization loss when performing online back-propagation. We omit the corresponding update expressions.

Limitations of BDIA-transformer: To our best knowledge, all existing reversible DNN models do not need to store any lightweight side information in the forward process. The primary objective of most existing works is to design reversible DNN models (e.g., RevViT) that produce comparable validation performance as the original counterparts. From the perspective of memory reduction, existing reversible DNN models are more efficient than
BDIA-transformer.

On the other hand, BDIA-transformer is designed to not only save training memory but also improve the generalization performance via model regularization. Our method is the first of its kind that attempts to maintain the transformer architecture in the inference procedure.

###### Remark 2 .

## 5 Experiments

We evaluated BDIA-transformer for three different tasks: (1) training BDIA-ViT for image classification; (2) training BDIA-transformer for language translation; (3) training BDIA-GPT2 for text prediction. The hyper-parameter $l$ for quantization in all three tasks was set to $l=9$. Three open-source repositories(^2^22[https://github.com/kentaroy47/vision-transformers-cifar10](https://github.com/kentaroy47/vision-transformers-cifar10) [https://debuggercafe.com/language-translation-using-pytorch-transformer/](https://debuggercafe.com/language-translation-using-pytorch-transformer/) [https://github.com/karpathy/nanoGPT](https://github.com/karpathy/nanoGPT)) were used in the experiments.

In brief, the obtained results from the first two tasks indicate that BDIA-ViT and BDIA-transformer significantly improves the validation performance of the original counterparts due to the model regularization effect of the random $\{\gamma_{k}\}_{k=1}^{K-1}$ variables. The results for the third task demonstrate that BDIA-GPT2 alleviates the overfitting issue of GPT2 significantly. RevViT from [] was also tested for the first task. It was found that RevViT does not always improve the performance of ViT.

Figure: Figure 3: Performance comparison of ViT, RevViT [], and BDIA-ViT for image classification over CIFAR10 and CIFAR100. $\{\gamma_{k}\}_{k=1}^{K-1}$ in the training procedure of BDIA-ViT were drawn from $\{\pm 0.5\}$ per training sample.
Refer to caption: x3.png

### 5.1 On training BDIA-ViT

In this experiment, we trained BDIA-ViT with $K{=}6$ transformer blocks on CIFAR10 and CIFAR100. The performance of ViT and RevViT [] was also evaluated to facilitate comparison. When implementing BDIA-ViT, the $\{\gamma_{k}\}_{k=1}^{K-1}$ parameters were drawn from $\{-0.5,0.5\}$ with equal probability per training sample. In addition, we utilized the SET-Adam optimizer [] in the training process with the configuration $(\eta_{0},\beta_{1},\beta_{2},\epsilon)=(1e{-}4,0.9,0.999,1e{-}18)$, where $\eta_{0}$ denotes the initial learning rate. The remaining training setups follow directly from the original open source (i.e., the first github link of footnote 2). Three experimental repetitions were performed for each training setup to mitigate the effect of the random seed.

**Table 1: Validation accuracy and peak memory consumption for training three models over CIFAR10 and CIFAR100. $\{\gamma_{k}\}_{k=1}^{K-1}$ in the training procedure of BDIA-ViT were drawn from $\{\pm 0.5\}$ per training sample. The peak memory includes both the model parameters and the training states for a batchsize of 128.**
|  | RevViT [] | ViT | BDIA-ViT |  |  |  |
| --- | --- | --- | --- | --- | --- | --- |
|  | val. acc. | $\begin{array}[]{c}\textrm{peak memory}\end{array}$ | val. acc. | $\begin{array}[]{c}\textrm{peak memory}\end{array}$ | val. acc. | $\begin{array}[]{c}\textrm{peak memory}\end{array}$ |
| CIFAR10 | 86.22$\pm$0.42 | 572.7MB | 88.15$\pm$0.55 | 1570.6MB | 89.10$\pm$0.38 | 693.4MB |
| CIFAR100 | 61.89$\pm$0.31 | 572.7MB | 61.86$\pm$0.47 | 1570.6MB | 66.09$\pm$0.80 | 693.4MB |

Performance comparison: Table 1 summarizes the obtained validation accuracy and the peak memory usages. It is clear that BDIA-ViT produces significantly higher validation accuracy than ViT and RevViT for both CIFAR10 and CIFAR100. Fig. 3 further visualizes the training and validation curves. It is seen that even though the training loss of BDIA-ViT is higher than the other two models across all the epochs, the validation accuracy improves remarkably in the end of training. This indicates that training an ensemble of ODE solvers parameterized (see Subsection 4.2) by $\{\gamma_{k}\}_{k=1}^{K-1}$ indeed regularizes BDIA-transformer properly.

In contrast, RevViT yields either inferior or comparable validation performance to ViT (see Table 1). This may be because RevViT modifies the architectures of ViT considerably to enable reversibility and therefore implicitly imposes uncontrolled regularization on the original transformer model. On the other hand, the regularisation introduced in BDIA-transformer is motivated by averaging consecutive integration approximations of an ODE in the original transformer, which leads to consistent performance gain across different datasets.

It is also clear from Table 1 that RevViT is most memory-efficient. BDIA-ViT needs slightly more memory than RevViT because it needs to store the lightweight side information $\{\boldsymbol{s}_{k}\}_{k=0}^{3}$ of the first 4 transformer blocks. BDIA-ViT improves the validation performance of RevViT at the cost of additional memory for storing the side information. We can also conclude from the table that online back-propagation does indeed significantly reduce the training memory.

Ablation study regarding the $\{\gamma_{k}\}_{k=1}^{K-1}$ parameters: As we mentioned in Remark 1, if reversibility is not of primary concern, different values for $\{\gamma_{k}\}_{k=1}^{K-1}$ can be employed for model regularisation. We conducted ablation study by evaluating the impact of $\gamma_{k}\in\{0.0,\pm 0.25,\pm 0.5,\pm 0.6\}$, $k=1,\ldots,K-1$, on the validation performance of BDIA-transformer. For doing so, both the quantization and online back-propagation operations were turned off in BDIA-transformer. We emphasize that the non-zero $\gamma_{k}$ values were only utilized in the training process. At the inference stage, $\mathbb{E}[\gamma_{k}]=0$ was used for computing the validation accuracy.

**Table 2: Impact of the $\{\gamma_{k}\}_{k=1}^{K-1}$ parameter in BDIA-ViT (w.o. quantization and w.o. online back-propagation) on the validation accuracy in percentage.**
| $\{\gamma_{k}\}_{k=1}^{K-1}$ | 0.0 | $\{\pm 0.25\}$ | $\{\pm 0.5\}$ | $\{\pm 0.6\}$ |
| --- | --- | --- | --- | --- |
| CIFAR10 | 88.15$\pm$0.55 | 88.79=3$\pm$0.29 | 89.12$\pm$0.22 | 88.89$\pm$0.15 |

Table 2 summarizes the obtained validation accuracy for different setups of the $\{\gamma_{k}\}_{k=1}^{K-1}$ parameters. It is clear that when $\{|\gamma_{k}|>0\}_{k=1}^{K-1}$ , the performance of BDIA-transformer improves considerably in comparison to that of the conventional ViT (i.e., corresponding to BDIA-ViT with $\{\gamma_{k}=0.0\}_{k=1}^{K-1}$). The setup of $\{\gamma_{k}=\pm 0.5\}_{k=1}^{K-1}$ performs the best. In general, the larger the magnitude of the $\{\gamma_{k}\}_{k=1}^{K-1}$ parameters in $[0,0.6]$, the slower the training speed (see Fig. 3). In practice, one can tune the magnitude of the $\{\gamma_{k}\}_{k=1}^{K-1}$ parameters within $[0,0.6]$ to balance the trade-off between the training speed and the validation performance.

As we mentioned in Remark 2, side information can also be introduced for the setups of $\{\gamma_{k}=\pm 0.25\}_{k=1}^{K-1}$ to make BDIA-transformer reversible. However, side information of 2 bits per activation value per transformer block is required instead of 1 bit for $\{\gamma_{k}=\pm 0.5\}_{k=1}^{K-1}$. On the other hand, it is complicated to enforce exact bit-level reversibility for $\{\gamma_{k}=\pm 0.6\}_{k=1}^{K-1}$.

Figure: Figure 4: Performance comparison for English to French translation. $\{\gamma_{k}\}_{k=1}^{K-1}$ in the training procedure of BDIA-ViT were randomly drawn from $\{\pm 0.5\}$ per training sample.
Refer to caption: x4.png

### 5.2 On training BDIA-transformer for English-French translation

Language translation is a classical natural language processing (NLP) task where the transformer shows a large performance gain over other DNN models. We adopted an existing open-source repository (i.e., the second github link of footnote 2) in our experiment. The tested BDIA-transformer has six transformer blocks in both the encoder and decoder, respectively. The BDIA update expressions were implemented in both the encoder and decoder. The performance of the conventional transformer was tested as a reference.

Fig. 4 visualizes the training and validation losses of the two tested models. It is clear that the obtained curves in Fig. 4 exhibit similar properties to those in Fig. 3. That is, even though the training loss of BDIA-transformer is higher than the conventional transformer across all the epochs, its validation loss is significantly lower after epoch 200.

### 5.3 On training BDIA-GPT2

In this experiment, we consider training BDIA-GPT2 over the dataset of openwebtext by utilizing the third github link of footnote 2. Our primary objective for this task is to find out if BDIA can help to alleviate the over-fitting issue of GPT2 (we omit the phase “nano" for simplicity) for a very small training dataset. In doing so, we only took a small (i.e., 0.05%) subset from the entire dataset when training the model. The performance of GPT2 was evaluated as a reference. Both models have 12 transformer blocks.

Fig. 5 shows the training and validation curves for the two models. Similarly to Fig.3-4, BDIA-GPT2 exhibits a slower training speed than GPT2. Considering the validation performance, the two validation curves for the two models exhibit the over-fitting issue. GPT2 produces the lower validation loss in the middle of the training. However, at the end of training, the validation loss of BDIA-GPT2 is significantly lower than that of GPT2. This demonstrates that BDIA indeed alleviates the over-fitting issue of GPT2 for a very small training dataset.

Figure: Figure 5: Performance comparison when training GPT2. $\{\gamma_{k}\}_{k=1}^{K-1}$ in the training procedure of BDIA-ViT were randomly drawn from $\{\pm 0.5\}$ per training sample.
Refer to caption: x5.png

## 6 Conclusions

In this work, we have proposed the BDIA training technique to assist the training procedure of transformers. Firstly, each transformer block is taken as the Euler integration approximation for solving an ODE. The BDIA technique is then applied to average every two consecutive integration approximations in a transformer as a regularizer via a set of random variables $\{\gamma_{k}\}_{k=1}^{K-1}$, one variable per transformer block. Exact bit-level reversibility for lossless online back-propagation can be achieved for BDIA-transformer by performing activation quantization and storing only lightweight side information. The side information accounts for the binary quantization loss incurred by the special setup of $\{\gamma_{k}=\pm 0.5\}_{k=1}^{K-1}$. In the inference procedure, the expectation $\mathbb{E}[\gamma_{k}]=0$ is taken to replace $\gamma_{k}$ in the update expression of BDIA-transformer, which reduces the update expression to that of the conventional transformer up to activation quantization.

Experiments on image classification and language translation show that BDIA-transformer produces significantly better validation performance than the corresponding baseline transformers while, at the same time, reducing training memory by performing online back-propagation. Experiments on text prediction indicate that BDIA-GPT2 alleviates the over-fitting issue of GPT2 over a very small training dataset significantly. In comparison, the experiments with RevVit demonstrate that it yields either inferior or comparable validation performance to that of ViT.

## Peer Reviews

confidence: 2
contribution: 3
initial_rating: 5
presentation: 3
questions: * What are the failure modes of BDIA-transformers? In which scenarios can the BDIA-transformers underperform compared to other transformers?

* Could you also report the comparison of BDIA-transformers and other architectures in terms of training time and computational cost?

* Minor: In Figure 3, the y-axis label should be “training loss”
soundness: 3
strengths: * Clarity: The authors clearly define the problem, related work, and their proposed method with well-written equations.

* Significance: The paper proposes a unique approach to achieve bit-level reversibility in transformers without architectural changes, leveraging techniques from ODE solvers and quantization. The introduction of the binary random variables also serves as a good regularization strategy.
summary: The paper proposes BDIA-transformer, a novel type of reversible transformer based on the bidirectional integration approximation (BDIA). A random hyperparameter $\gamma$ is introduced per transformer block per training sample to regularize the models. The paper further performs activation quantization to allow for exact bit-level reversibility of BDIA-transformers. Empirical results show that the BDIA technique outperforms baseline transformers and reduces training memory for image classification and language translation.
weaknesses: * Lack of experiments: The paper compares BDIA-transformers with standard transformers and RevViT. It would benefit from a broader evaluation against other reversible architectures or quantization methods.

* The reliance on storing lightweight side information for exact reversibility might reduce its practical applicability as depth increases. It would be nice if this trade-off between memory efficiency and information storage is more analyzed.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 2
questions: - Can you provide more intuition for the choice of $\gamma$ values? Why are these values chosen, and what do they represent? Further, do you have any theoretical insights into the choice of these values?
- Can you clarify, more explicitly, how you arrived at 1 bit for $\gamma \in \{0.5, -0.5\}$ and 2 bits for $\gamma \in \{0.25, -0.25\}$? How does this relate to the $\mathbf{s}_{k-1}[m]$ variable from equation (20)?
- In section 5.1, you mention hat you're training with $K=6$ transformer blocks. Later, when talking about the memory overhead, however, you denote the side information as $\{s_k\}_{k=0}^3$. Can you clarify this? Why is the side information only stored for the first 4 blocks?
soundness: 3
strengths: - The background on the BDIA technique and its application in diffusion inversion is well-written and provides a natural motivation for its application in transformers.
- The authors provide a clear motivation for their work and the necessity of reversible transformer architectures.
- The application of this technique to transformers is novel and well-motivated.
    - More specifically, the connection between diffusion inversion and neural ODEs is an interesting one and provides a unique perspective on the problem.
- Experimental results demonstrate the effectiveness of the proposed BDIA-transformer on various tasks.
summary: This paper presents BDIA-transformer, a novel approach to creating reversible transformers that maintain their original architecture during inference. The authors incorporate the recently proposed bidirectional integration approximation (BDIA) technique, originally proposed for diffusion inversion, into the transformer architecture training process.

During the training process, a hyper-parameter $\gamma \in \{0.5, -0.5\}$ is randomly selected for each sample and transformer block to average consecutive integrations. This approach effectively trains an ensemble of ODE solvers. At inference time, the expectation $E(\gamma) = 0$ is used, reducing the model to a standard transformer with quantized activations.

The authors observed that the BDIA update expression is only theoretically reversible when using floating-point arithmetic, leading to error accumulation, especially for deep networks. To address this issue and enable lossless online backpropagation, they apply activation quantization to achieve exact bit-level reversibility.

Experimental results on various tasks, including image classification, language translation, and text prediction, demonstrate that BDIA-transformer (1) uses significantly lower overall memory during training (compared to ViT) and (2) acts as a regularizer, improving validation accuracy (over both RevViT and ViT) and reducing overfitting on small datasets.
weaknesses: My main concerns with the paper are with the presentation and clarity of the methodology. Specifically:

- The paper begins (both in the abstract and introduction) with heavily references to the $\gamma$. However, the significance or meaning of this parameter is not immediately clear to the reader.
- The choices for specific values of $\gamma$ are not well-motivated. Further, for each $\gamma$ choice, seemingly different amounts of side-information are required, but these details are also not explained well.
- Similar to the first point, the reason for the existence of the activation quantization is not clear until much later in the paper. This is also something that should be explained much earlier in the paper (e.g., in the introduction).

---

confidence: 2
contribution: 3
initial_rating: 3
presentation: 1
questions: - How does the extra computation time for a BDIA transformer compare to that of a regular transformer? In what situations does this additional overhead make sense? Could the authors also report the wall-clock performance of their method?

- Could the authors explain why they view a sequence of transformer blocks as a single ODE solver? This idea seems to fit better with diffusion models, as proposed by Zhang et al. 2023, but it’s not as clear for transformers.

- How well does the BDIA-transformer handle larger datasets or bigger models? Would the memory and speed trade-offs still hold up in these cases?
soundness: 2
strengths: - The paper addresses the “memory wall” challenge, a critical issue in deep learning in which training large models requires extensive memory to store intermediate activations alongside the model itself. Authors explore reversible deep learning to tackle this issue. Reversible architectures allow backpropagation without storing or with minimal storage of intermediate activations, thus offering significant memory savings. The authors propose a novel technique by treating each transformer block as an Ordinary Differential Equation (ODE) solver and applying the Bidirectional Integration Approximation (BDIA) to these blocks. 

- Experimental results show that BDIA-transformers improve performance across various tasks, including image classification and language translation. Authors claim that performance gains and memory efficiency are gained by implementing this method on Vision Transformers.

- This approach has the potential to enable the training of larger, more complex models when memory is the bottleneck, but we have very strong computational power.
summary: The paper introduces BDIA-transformers, a novel approach to reversible transformers designed to reduce memory consumption during training without altering the architecture during inference. The approach leverages the Bidirectional Integration Approximation (BDIA) technique, which treats each transformer block as an Euler integration approximation for solving ordinary differential equations (ODEs). The experimental results show that BDIA-transformers outperform standard transformers in image classification and language translation tasks while reducing training memory requirements. For text prediction, BDIA-GPT2 prevents overfitting when trained on small datasets.
weaknesses: - The paper is hard to follow, especially in the preliminary and method sections. It doesn’t clearly explain key ideas, like why transformers can be viewed as ODE solvers or how BDIA transformation is applied, which makes it confusing.

- The related work section mainly lists past studies without explaining how the field has developed or why this approach is needed. This makes it hard to see where this paper fits in with previous work. 

- Some arguments supporting the method are vague. For instance, in the subsection "on similarity to dropout technique". 

Overall, a clearer, simpler structure and explanations would make the paper easier to understand, especially for readers unfamiliar with reversible deep learning.

- The paper attributes BDIA-transformer's improved performance to the regularization effect of random variables, which is said to work similarly to dropout. However, it’s not clear if this is the only reason for the performance gains. The authors haven’t compared BDIA against other standard regularization methods on the baseline transformer, leaving open the chance that similar improvements could be achieved without the added complexity of BDIA.

- For language models, benchmarks and perplexity provide more informative insights into model quality. Could the authors test their methods against GPT-2 using additional metrics?

---

confidence: 2
contribution: 2
initial_rating: 5
presentation: 4
questions: Could the authors elaborate on the three weaknesses ?
soundness: 2
strengths: 1. The paper is well-written and the algorithm is pretty easy to follow.
2. The experiments, while being rather small scale, do demonstrate a surprising regularization effect, able to tackle the overfitting issue of transformers.
summary: The presented paper proposes a new training algorithm for transformers enabling reversibility up to a fixed precision level. While the paper itself focuses on transformers, the method seems to be applicable to any residual architecture. The method itself enables exact reversibility up to a given precision level, without architecture modification, and is thus broadly applicable. The authors introduces a regularizing parameter $\gamma$ which substantially alleviates overfitting issues in Transformers.
weaknesses: 1. The parameters $\gamma_k$ seems to already exist in the original BDIA paper and it is thus not clear what is the novelty of this paper on the matter.
2. The exact reversibility seems to require the quantization step but there is no ablation on the precision level $l$.
3. The small scale experiments does demonstrate a surprising regularization effect. However, the paper seems to seek reversibility, which is a feature usually used to scale up the model size given a fixed compute setup. Therefore, it is quite strange to focus on small scale experiments.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: -	In figure 1, and in line 287, how did the authors integrate $\gamma$ into standard transformers?
-	In figure 2, the authors should the reconstruction errors w.r.t. the proposed method using quantization and side information. Otherwise, it is not clear the effectiveness of these tricks.
-	The authors should compare experimentally the proposed methods against vanilla transformers applied with dropout.
-	What dataset did the authors use in the machine translation experiments.
-	Although the authors show the memory gains, they should show the convergences in terms of wall-clock time to see better the computational complexity introduced by the proposed method.
soundness: 1
strengths: - The paper is generally well-written.
- The paper addresses an important and timely problem: reducing the memory consumption during the training of transformers, which is particularly relevant given the current widespread use of transformer models.
- The proposed idea is compelling as it retains the original architecture of transformers. This stands in contrast to existing approaches that typically involve modifications to the transformer architecture.
summary: The paper proposes a novel type of reversible transformers with the aim to reduce the memory during training. To this end, this work treats each transformer block as the Euler integration approximation in a manner similar to Neural ODEs. There are two main contributions. Firstly, the authors borrow a technique from recent works on diffusion inversion for round-trip image editing, which involves bidirectional integration approximation. This approximation introduces a hyperparameter $\gamma$. The authors propose selecting $\gamma$ randomly either -0.5 or 0.5 for each training sample and training block. Consequently, the training can be viewed as an ensemble of ODE solvers. This regularization led to observed improvements on validation data. Secondly, to ensure reversibility, the authors propose performing activation quantization while storing side information. This approach is validated on small datasets involving image classification, machine translation, and language modeling.
weaknesses: -	The reproducibility is low as there is no source codes or pseudo codes or detailed algorithms.
-	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. Notably, in equation (4), the authors treat the combined attention and feed-forward network modules as a residual term, resulting in derivations similar to those found in NeuralODEs with ResNets. However, these modules are key differentiators in transformer architectures compared to other models.
-	The experiments mainly consider small datasets or relies on toy examples for transformers.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: + There is a tipo at line 169 (written $\\{-1/2, -1/2\\}$ instead of $\\{-1/2,1/2\\}$)
soundness: 4
strengths: + The paper is well-written, the ideas are clearly presented
+ Several insights on how the method works and relates to previous work are provided such as its similarity to dropout
+ The experiments effectively support the claims made in the article, such as the regularization effect of the method
summary: The article derives a new method to make a transformer exactly reversible, based on a method initially dedicated to diffusion models and on the quantization of the activations. The proposed method also seeks to regularize the training of the model. Experiments are conducted on a variety of benchmarks with different transformer-based models, that is, vision-transformer, vanilla transformer and GPT-2. The proposed method achieves SOTA results on a vision task, while also preventing overfitting.
weaknesses: + The related work section is very short and surprisingly focuses on quantization, while quantization appears more in this work as a convenient trick to make information representation simpler rather than the primary topic of the article. Notably, the presentation of other works on reversible networks is missing (in this section)
+ The whole section on quantization (section 4.3) appears a bit messy. For instance, there is a parameter $l$ which is never really discussed, making it seem quite arbitrary. Similarly, the information storage vector $s_k$ is defined in a technical manner, and the section would benefit from providing more intuitive insights into its function.