On Exact Bit-level Reversible Transformers Without Changing Architecture
Authors: Guoqiang Zhang
JP Lewis
W. Bastiaan Kleijn
Venue: ICML 2025 poster (2025)
Keywords: transformer; reversability; BDIA; quantization
Primary area: deep_learning->algorithms
TL;DR: A new type of reversible transformer for online back-propagation that improves generalization performance and becomes the standard transformer in inference procedure

Abstract: In this work we present the BDIA-transformer, which is an exact bit-level reversible transformer 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$ in BDIA-transformer randomly take one of the two values $\{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 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 $\mathbb{E}(\gamma)=0$ is taken to make the resulting architecture identical to transformer up to activation quantization. Our experiments in  natural language generation, image classification, and language translation show that BDIA-transformers outperform their conventional counterparts significantly in terms of validation performance while also requiring considerably less training memory. Thanks to the regularizing effect of the ensemble, the BDIA-transformer is particularly suitable for fine-tuning with limited data. Source-code can be found via \href{https://github.com/guoqiang-zhang-x/BDIA-Transformer}{this link}.

On Exact Bit-level Reversible Transformers Without Changing Architecture
Guoqiang Zhang'
J.P. Lewis?
W. Bastiaan Kleijn
Abstract
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 vari- ous tasks. Examples are large language models (LLLM) such GPT-4 (Achiam et al., 2023) and Llama 2 (Tou- vron et al., 2023), which can, for example, have informative and friendly conversations with humans, solve mathemat- ical 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 accommodate both the DNN model and the interme- diate activations of the input data-batch in back- propagation (Gholami et al., 2024), which is referred to as memory wall in the literature.
In this work we present the BDIA-transformerr, which is an exact bit-level reversible transformer that uses unchanged standard architecture for inference. The basic idea is to first treat each transformer block as the Euler integration approx- imation for solving an ordinary differential equa- tion (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 y in tDaA-transformer randomly take one of the two values (0.5. ) per training sample per transformer block for averaging every two con- secutive integration approximations, As result, BDIA-transformer can be viewed as training an ensemble of ODE solvers parameterized by set of binary random variables, which regularizes the model and results in improved validation accu- racy. Lightweight side information 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 ex- pectation (B))) = 0 is taken to make the resulting architecture identical to transformer up to acti- vation quantization. Our experiments in natural language generation, image classification, and lan- guage translation show that EDDIA-transformet outperform their conventional counterparts signif- icantly in terms of validation performance while also requiring considerably less training memory. Thanks to the regularizing effect of the ensemble, the BDIA- transformer is particularly suitable for fine-tuning with limited data. Source-code can be found via this link.
One promising technique to alleviate the issue of memory wall to design and train reversible DNNs (Dinh al., 2014; 2016; Behrmann et al., 2019; Mangalam et al., 2023). 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 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 (Dinh et al., 2014), the methods of (Rezende & Mohamed, 2015), Real NVP (Dinh et al., 2016), and Inverse Autoregressive Flow (Kingma et al., 2016) differentiable transformations that were aimed at facilitating generative modeling. For overview of such early normalizing flows see (Kobyzev et al., 2020). Inspired by these works, a range of reversible residual models were proposed subsequently, including RevNet (Gomez al., 2017), Glow (Kingma & Dhariwal, 2018), i-RevNet (Jacobsen al., 2018), i-ResNet (Behrmann et al., 2019), layerwise inversion (Hascoet et al., 2019), Fourier- transformation based CNN inversion (Finzi et al., 2019), and Mintnet (Song et al., 2019). Another line of research enforces reversibility in deep learning from the perspec-
'Department of Computer Science University of Exeter, UK "NVDIA, USA School of Engineering and Computer Science, Victoria University Wellington, New Zealand. Correspondence to: Guoqiang Zhang <guoqiang.x zhzaangg.....agggg...o. @ gmail.com>.
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

On Exact Bit-level Reversible Transformers Without Changing Architecture
tive of ordinary differential equations (ODEs), which in- cludes FFJORD (Grathwohl et al., 2018), leapfrog networks (Chang et al., 2017), momentum residual networks (Sander et al., 2021), and neural ODE inversion (Stam, 2022).
mance over that of the corresponding baseline transformers and simultaneously reduces training memory significantly. The improved performance results from the model regu- larization imposed by a set of random variables. Our empirical study also indicates that Rev from (Mangalam et al., 2023) produces either inferior or comparable valida- tion performance to that of its original counterparts.
Recently, the research on reversibility has moved to other types of neural networks. The authors in (Mangalam et al., 2023; Zhu & Mangalam, 2023) proposed reversible vision transformers (referred to as RevViT) due to the popular- ity of LLMs. (Wallace et al., 2023) utilized reversible diffusion sampling method 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 architec- tures considerably to enable reversibility.
2. Related Works
In recent years, various quantization strategies (Yang et al., 2019; Wu et al., 2018; Wang al., 2018) have been pro- posed in the training and/or inference processes of DNN models on low-precision devices. For instance, the work (Wu et al., 2018) successfully performed quantization on DNN weights, activations, gradients and errors in the train- ing and inference processes and obtained promising results. The recent work (Ma, 2024) demonstrated that L.MSs with quantization of 1.58 bits per model parameter exhibit com- parable 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- reversibility in training BDIA-transforme transformers.
new type of reversible transformer that uses unchanged, standard ar- chitecture for the inference procedure. It is based on the bidi- rectional integration approximation (BDIA), which was re- cently proposed in (Zhang et al., 2023a) to enable diffusion inversion for round-trip image editing. To be able to incor- porate BDIA into transformers for online back- propaaaaaag,,nn, we follow the common practice of treating each transformer block as Euler integration approximation for solving ordinary differential equation (ODE).
In this paper, we propose
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 (Ma, 2024), advancing the research frontier of reversible DNN models has become of great interest.
We make two main contributions in this work. Firstly, we propose BDIA- tranfforrrfsssss by introducing a random hyper- parameter y E (-0.5, ) per transformer block per train- ing sample to regularize the DNN models for improvement of validation performance. Each 7 parameter intends to aver- age 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 E[>] = 0 is uti- lized, which reduces BDIA-trranfoommers to conventional transformers. As a result, our method is good candidate for fine-tuning existing transformer-base models such as LLMs..
In addition to reversible DNNs, other nemory-saving tech- niques have been proposed in recent years, such as gradient checkpointing (Feng & Huang, 2021), and activation of- floading (Wu et al., 2024). At least in principle, the above techniques can be combined with the BDIA algorithm to save training memory for a broad class of residual-type neural architectures.
Secondly, we perform activation quantization to allow for exact bit-level reversibility of BDIA-transtormert Note that the special setup of the y values in the subset (-0.5, 0.5) when performing activation quantization leads to a 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.
3. Preliminary
Neural networks as ODEs: (Chen et al., 2018) highlighted the interpretation that passing a hidden state across layers that add a correction to that state can be viewed as Euler inte- gration. While that paper identified architectures including residual nets and normalizing flows as following this pat- tern, it is equally true of diffusion models and transformers. This interpretation emphasizes the importance of consider- ing more accurate integration schemes (Karras et al., 2022). The need for improved integration is particularly apparent in round-trip image editing in diffusion models, where sig- nificant integration error will result in unintended visible
Experimental results for natural language generation (NLG), image elassification, and language translation show that the BDIA technique significantly improves the validation perfor-
See Subsection 5.1 for full fine-tuning and Appendix E for LoR A based fine-tuning.
2

---

On Exact Bit-level Reversible Transformers Without Changing Architecture
alterations to the image.
inversion for round-trip image editing. By reformulating (3), 2;-1 can be easily computed in terms of (zi, Z2+1). We note that due to the nature of the floating-point datatype, there might be error accumulation in round-trip image re- construction, 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 (Zhang et al., 2023a) 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.
Difftssonn sampling via solving ODE: Recently, the work (Zhang et al., 2023a) 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 2T at time step T, the diffusion- iaffusiio-bbasd sampling process for generating realisic images ze at time e > 0 can be realized by solving a probability ordinary differential equation (ODE)
dz d(z,()))) ()tt
One main difference between diffusion inversion and re- versible transformers is that no gradient needs to be back- propagated in diffusion inversion for updating the DNN model. As result, even if there is error accumulation in dif- fusion inversion, it is less severe than in reversible transform- ers where error accumulation in online back- pack-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.
over the time interval t E [T. The gradient vector d(z,t) t) includes the output of pre-trained DNN model with (z,, as its input. The common practice for solving the above ODE is to first discretize the continuous time interval [T", el properly into a set of timesteps {4//i = .....,N } with to = T and ty = and then perform certain integration ap- proximation per small time-interval sequentially to compute the final diffusion state N Z..
BDIA: Suppose we would like to estimate the next diffu- sion state 2++1 2t,41 by solving (1) based on the re- cent information (z,, f.) and (zi-1,t--1), where 2j 21, for j i i, The basic idea of BDIA is to com- pute z;+1 by performing both the forward integration ap-
4. Exact Bit-Level Reversible Transformers via BDIA
proximation A(t; + t++1lz,) [= S.i d(z, t)dt) and the backward integration approximation A(4; - 4-1122)) ( d(z,()d) conditioned on z;. One popular method for implementing Att; - (+,1 and A(t, 4-1122) 2;) in the literature of diffusion models is by employ- ing the DDIM update expression (see (Song et al., 2021; Zhang et al., 2023b;a) for details). With the above two integration approximations, 2++1 can be expressed as
In this section, we first briefly review the transformer update expressions from the ODE viewpoint. We then propose the that enables exact bit-level reversibility with activation quantization. Specially, we will demonstrate how each of the two values { (-0.5, 0.5] averages consec- utive integration approximations. The training of BDIA- transformer can then be interpreted as employing different ODE solvers for different training samples in random manner. In addition, explain why additional lightweight side- inde-informatin is required to be stored to account for the binary quantization loss in online back propapropa.
Zi+1 =2 -1 - 7) 2,) - AA(4 + 4-122)
214n d(z,))d
+ A(li + أعااا
4.1. Revisiting transformer update expression
2 1 d(ztt)
A typical transformer block consists of the attention function (denoted as f-)) and the function of feed-forward network (FFN) (denoted as g()).. of which the trainable parameters are generally different for different block indices. Accord- ingly, the output x++11 of the kth transformer block can be mathematically represented in terms of the input ag as
- 1 + (1 - 7) Zii - -A - 1-1121)
+ At, + 1 124);
(3)
where y= = (0, 1] averages Al4: + (211) and (z,-1- - 2.) for the time-slot [t,-1, 1.]. The minus sign in front of the two quantities are due to the reverse integration direction. The quantity - - 2;) is the previously computed in- tegration approximation for Jh d(z, t)tt. We note that negative 7 values are not applicable to the diffusion models considered in (Zhang et al., 2023a).
x++1 = D& + (x() + gh g + fx()))
(4)
(mw)
where we use h2 (z2) to denote the overall residual quan- tity that includes both the attention and FFN functions. For simplicity, we omit the pre-normalisation operations
On reversibility of BDIA: The update expression (3) is care- fully designed in (Zhang et al., 2023a) to enable diffusion
3

---

On Exact Bit-level Reversible Transformers Without Changing Architecture
in (4), since these in fact do not affect the design of BDIA- transformers later on.
Inforence by different ODE solvers parantertiod by
7 1 (1 2/2x + (1 1 ( (1
It is well known from the literature (Chen et al., 2018) that the kth forward step in (4) can be roughly viewed as the Euler integration approximation of an ODE at timestep 42:
68
68
Alty-Hul (Hl1am)))
BDIA-VI
2h+1 Fh (=(()))=0000 = + dla (1x+1 (b) (5)
0.5
-0.4
-0.2
-0.1
0.3
0.4
/ 1x+1
d(x, t)dt,
Figure I. Validation performance of diffcrent ODE solvers param- eterized by a single y parameter after training VYT and BDIA-ViT over CIFARIO. See Subsection 5.2 on how VIT and BDIA- VIT were trained. Each ODE solver in the inference procedure is re- alized by selecting "y 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.
where d(zs, 1x) denotes the gradient vector with (ate, 4x) as the input. Both d(æ;, (x) and (f++1 - 4)) are implic- itly learned via the composite function (xx)),), which is alternatively denoted Att + (xx)
As explained later on, will introduce BDIA into the up- date expression of (5). Note that (5) is a general update expression that not only includes the transformer but also is noted in
In the inference stage, E(y) = 0 is taken to replace y in (10), which leads to simpler update expression that only involves (Ig,, xx++1):
ResNet (He et al., 2015) and its variants. (Dupont et al., 2019) that ResNet can represent more func- tions than neural ODEs. Despite the less powerfulness of the ODE framework, we employ the framework to facilitate the design of reversible transformer and obtain promising empirical results.
0k+1 = By. + hxx()))
(11)
which is identical to the original transformer update expres- sion (4)-(5).
4.2. BDIA-transforman without quantization
The above analysis suggests that the BDIA training tech- nique is good candidate for fine-tuning transformer-based models because BDIA retains the standard architecture in the inference procedure. Fine-tuning pre-trained model is typically performed on small or intermediate-sio dataset, which is prone to overfitting. Subsection 5. 1.1 and Appendix E demonstrate successfully applying BDIA to fine-tune GPT2 medium using small dataset.
In this subsection, first derive the update expressions of AiDI--transfomme without quantization as an extension of (5), and then study the impact of the y values in {0.5, -0.5). When k = 0, Φ can be computed by following (5) as
Wo + ho(x)) xo + Ato t (z)).
The update expression for xk+11 in the training process, K 122 k 2 can be obtained by utilizing (2)-(3). Based on (5), we let
Impact of y parameter: We now study the impact of the 7 parameter in (10). When y 0.5, it follows from (2)-(3) and (7)-(8) that the two integrations f d(x,, r)dr and f d(x,, T)dr of the transformer ODE are approximated as
(tx + t--1 (x)) = (xx)
(7)
A(t: t++1 (xx) = (xxx)))
By combining (7)-(8 and (2)-(3) with i = k; the update expression 1x+1, K 2 t N 1, can be represented as
f d(x,, r) zo 0.5(xxx +02)-mm---)+
(12)
2x+1 2 24-1 [ 2 + (mx) (9)
pati [(
= -
(1 (mx),
(10)
d(æt, r)dr so hx(
(13)
where y is recommended to take values from (0.5, 0.53 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 ADIA-based diffusion inversion (Zhang et al., 2023a) where y has to be positive.
where the integration over [44-1, 4. is computed as the weighted average of two consecutive integration approxima- tions: (æt xk-1) and hu (2x).
When 7 = -0.5, the two integrations fi,, d(x,, T)dT and

---

On Exact Bit-level Reversible Transformers Without Changing Architecture
J a(æ, T)d7 are approximated differently, given by
x10
without quentzztion
St" d(x,, T)dT x (xx ( 1))
quantization side info
(14)
1.5
tk-1
(let) f 7)dT = 0.5hg (xs)) -0.5(20000000-- (15)
recon.
0.5
In this case, the integration over [bx: (x+1] is computed by averaging (xx - 1 and (xx)).
x0 x2
x3
x4
x5
x6
x7
X8
x9
x10
blockwise activations
Training via an ensemble of ODE solvers: We use 7 to denote the random variable for the kth transformer block taking values in {0.5.- 3 with equal probability. The above analysis of (12)-(15) implies that each training sam- ple goes through particular ODE solver determined by sampling set of K - 1 random variables (7*) k== which specifies the integration path from the bottom transformer block until the top one. Due to randomness, different train- ing samples will go through different ODE solvers. There are in total 2K different ODE solvers, where each one corresponds to unique integration path across the K trans- former blocks. Intuitively speaking. if we assume that each of the individual ODE solvers is well behaved (i.e., its train- ing loss decays to a small value), then a convex combination will also converge to a small value.
Figure 2. Demonstration of the accumulated reconstruction errors for two scenarios when training BDIA-GPT2 with 12 transformer blocks using the setup 7 (0.5, 00.5), k = -....... 11. The two scenarios are BDIA without quantization by following (16), and BDIA with quantization and side information by following (24). See also Fig. 6 for visualization of the reconstruction errors.
maintaining the original DNN architectures in the inference procedure. See Table 3 for ablation study of the impact of the {7w ] K-1 when training
parameter on the validation performance for image classification.
4.3. On exact bit-level reversibility of BDLA-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 enable exact bit-level reversibility. We note that lightweight side information is required to be stored per transformer block for lossless online back-propagatiinn
In principle, different ODE solvers can also be applied in the inference procedure. For instance, one can set y to be constant within [-0.5, 0.5] in (10) across all the transformer blocks for the same input. Fig. 1 demonstrates the vali- dation performance of those different ODE solvers in the inference procedure for both trained BDIA-ViT and ViT over CIFARIO. It is clear that the validation performance of BDIA-ViT is much more insensitive to the single 7 parame- ter than that of ViT. The reason for the sensitivity of VIT to y in Fig. is because only a single ODE solver is trained for VIT.
Limitation of the reversibility of (10): The update expres- sion (10) is only theoretically reversible. That is, æx-1 can be computed in terms of (Ix 2k+1) as
On similarity to dropout technique: From a high level point of view, the training of Bran--trasfoommer is similar to the conventional dropout technique (Srivastava et al., 2014) to a certain extent. Dropout essentially attempts to train an ensemble of subnetworks of the entire DNN model when minimizing the objective function while ADDA--tranforrer intends to train an ensemble of different ODE solvers. In the inference procedure, an average of the ensemble of sub- networks in dropout is utilized while in BDIA- trrnsorrroorrrr the expectation E[xx] = O is employed to replace "Ik- In the experiments we investigated the joint impact of dropout and BDIA for the image classification task, and obtained positive results. See Table 5 for details.
1
1- 7h or
1 1 ha (2x),
26
0k+1
(16)
7kk
7kk
Tkk
where we use " indicate that different transformer blocks have their respective random variables. In practice, the setup "/k E {0.5, -0.5). k = I.....K K 1, would lead to non- negligible error accumulation especially for very deep trans- former models. The factor = +2 in front of Øx+1 would amplify the error when k decreases from K - to 1, making the online back- pock-propagation unstable. Fig. 2 illustrates that the reconstruction error without quantization increases sig- nificantly when applying the online-back- bar--eaaaattoottttoonnn from the last transformer block back to the first.
Remark 4.1. If reversibility is not of concern, one can freely specify the values for the random variables (7x K -1 j-- in the training process for performance improvement as long as its distribution is symmetric around such that E/m]] = 0 for
AtDnn-rrasfrrmer with quantization: less online back proppprppag,tion we propose to perform acti- vation quantization. In particular, we use Qi(:] to denote
To allow for loss-
5

---

On Exact Bit-level Reversible Transformers Without Changing Architecture
quantization the bit-level precision of 2-1, given by
See Appendix C for the proof. The lossless reconstruc- tion (24) ensures that the computed gradients in the online back-propagation would not deviate from those in conven- tional back- propppppppgttonn which is desirable in very deep transformer models. In Fig. 2, the blue line shows the ex- act reconstruction of the intermediate activation values, as enabled by (24) during online back-propagation.
[ = round / 1|2
(17)
Upon introducing Q.1-) the new update expression for EDIA-transformer can be represented as
zo < (
(18)
21 = zo + (1)
(19)
Finally we briefly consider the inference procedure. Again replace 7 in (21) by E(*x)) 0. As result, the update expression (21) can be simplified to be
8k [m]] - {
1 if mod
k 221 (20)
otherwise
0s+11 (7 + 12
2 = 21 + h (xx))
k 221.
(25)
(10[(11%)) + (#+) k 2 1, (21)
where N% E {0.5, 0.5}. -..). and æz-1 denotes the mit element of -k-1- The math element s indicates if the integer value /m is odd or not.
The only difference of (25) w.r.t. the original transformer update expression (4) is that the quantization operation [ is performed for each activation output.
Lemma 4.2. Suppose {ax } k= are computed sequentially by following (18)-(21). Then the intermediate activation output wy has fixed-point precision 2-
Remark 4.5. The particular choice of * E {0.5, -0.5] is also motivated by exact bit-level reversibility. When y is power of two such as 0.5 2-11 the result of a product ? x is equivalent to shifting the mantissa of by one bit or decrementing its exponent, so at most exactly one bit is lost. This property is not true for arbitrary (non-power-of-2) values of y.
I& = ===21 [xw]
K K2 220 O.
(22)
We refer to the binary vector 8k-1 in (20) as the lightweight side information computed based on 03-1- The vector 8k-1 essentially captures the 1-bit quantization loss of [1 per element, and is crucial for recovering #g-1 exactly from (I), xk+1)) We now characterize the first quantization operation on the right hand side (RHS) of (21) with the following proposition:
Limitations of BDIA-trrnsformer: To our best knowledge. all existing reversible DNN models do not need to store any lightweight side information in the forward process. The pri- mary objective of most existing works is to design reversible DNN models (c.g., RevViT) that produce comparable vali- dation performance as the original counterparts. From the perspective of memory reduction, existing reversible DNN models are more efficient than BDIA Aall--ransformer.
Proposition 4.3. Suppose the conditions in Lemma 4.2 hold.
Thenfoor any z 1, we have
(1 + 12 8 ')1 = (I + 1 12 (23)
On the other hand, BDIA ADDAA-tranformer is designed 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.
The proof for Proposition 4.3 can be found in Appendix B. Equ. (23) indicates that the quantization operation has no effect on (x + 1 )) which naturally facilitates the exact reconstruction of - from (Ig,, xx-
5. Experiments
On reversibility of (21) by storing lightweight side in- formation: Suppose in each forward pass in the training process, all the side information 1 is stored in the memory. Then exact bit-level reversibility is guaranteed: Proposition 4.4. Suppose the conditions in Lemma 4.2 hold and the side information {8k-1) k= is available. Then {xx-1)K K-1 can be reconstructed exactly in the reverse order with the initial state-pair (IK-1, x:):
We evaluated BDIA-transformer for four different tasks: (I) fine-tuning GPT2 using the BDIA training technique for natural language generation (NLG); (2) training BDIA- VIT for image classification; (3) training Brrn--trnsformmer for language translation; (4) training BDIA-GPT2 for text generation. The hyper-parameter for quantization in all four tasks was set to / 9. Four open-source repositories were used in the experiments (see Table 7 in Appendix G). Due to the space constraint, we put the experimental results for the fourth task in Appendix F.
1
0k-1
0k+1 - -2-12-1
7
1
Qi (1-11 - + (1+
(24) In brief, the obtained results from the first three tasks in- dicate that the BDIA technique significantly improves the validation performance of the original counterparts due to
7
where k = K - 1
6

---

On Exact Bit-level Reversible Transformers Without Changing Architecture
the table that results for FT-BDIA are consistently better than those for FT. On the other hand, even though the LoRA technique is light-weight, its performance is promising.
Table 1. Performance for fine-tunning GPT-2 medium (M) on the E2E NLG Challenge. For all metrics, higher is better.
trainable param.
Remark 5.1. We have also investigated the joint impact of LoRA and BDIA when fine-tuning GPT2 M for NLG, and again we obtained positive results. See Appendix E for details. In brief, BDIA helps with improving the validation performance for LoRA based fine-tuning for the considered task. To enhance the impact of BDIA in the fine-tuning process, it is recommended to set a high value for the a parameter in LoRA.
Method
BLEU
NIST
MET
ROUGE-L
CIDE
LoRA
0.35M
69.0
8.70
46.4
71.3
2.51
FT
354,92M
66.8
8.52
46.3
70.3
2.37
FT-BDIA
354,92M
68.6
8.65
46.6
71.3
2.49
1.35
LORA
2.6
1.3
-FT-BDIA
2.4
1.25
5.2. On training BDIA-ViT
train loss
val loss
In this experiment, we trained BDIA-ViT with K=6 trans- former blocks on CIFAR1O and CIEAR 100 by using a single 2080 Ti GPU. The performance of ViT and Rev (Man- galam et al., 2023) was also evaluated to facilitate compari- son. When implementing BDIA-ViT the { param- eters were drawn from (00.5, 0.5) with equal probability per training sample. In addition, we utilized the SET- Adam optimizer (Zhang, 2024) in the training process with the configuration (No. 21,32,5, E) = (1e-4, 0.9, 0.999, 18-1 where 7/ denotes the initial learning rate. The dropout rate was set to 0.1 to reduce over-fitting. The remaining training setups follow directly from the original open source (i.c., the first github link of Table 7). Three experimental repetitions were performed for each training setup to mitigate the effect of the random seed.
2.2
1.15
1.1
|
FT-BDIA
1.8
1.05
0.5
1
1.5
2
2.5
0.5
1.5
2
2.5
iterations
x10
iterations
x10"
Figure 3. Performance comparison when fine-tuning GPT2 Medium for the E2E challenge. {7x (x L-1 in the BDIA training
procedure were randomly drawn from { + per training sam- ple. FT-BDIA refers to "fine-tuning via BDIA training technique" while FT refers to "fine-tune -unint directly".
the model regularization effect of the random {7H | vari- ables. The results for the 4th task demonstrate that BDIA- GPT2 alleviates the overfitting issue of GPT2 significantly for limited training data. VIT from (Mangalam et al., 2023) was also tested for the second task. It was found that Rev does not always improve the performance of VIT.
Table 2 summarizes the ob- tained validation accuracy, the peak memory usages, and the average training time per epoch. It is clear that BDIA- ViT produces significantly higher validation accuracy than ViT and Rev for both CIFARIO and CIFARION Fig. 4 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 train- ing. This indicates that training an ensemble ODE solvers parameterized (see Subsection 4.2) by { indeed reg- ularizes BDIA-transformer properly.
Performance comparison:
5.1. On fine-tuning GPT2 medium (M) for NLG
In the first experiment, we consider fully fine-tuning GPT2 medium (354.92M parameters in total) for the natural lan- guage generation (NLG) task by using the E2E dataset (Novikova al., 2017). We adopt the open-source for GPT2 (LoRA) in (Hu et al., 2021) (the second github link of Ta- ble 7) to evaluate the BDIA training technique. For com- parison, we also fine-tune GPT2 directly and via the LoRA technique with the default setup of (rank, a) = (4, 32). It is worth noting that the LoRA technique only needs to train a very small number of parameters.
K-
In contrast, RiT yields either inferior or comparable val- idation performance to ViT (Table 2). This may be because Rev modifies the architectures of ViT considerably to enable reversibility and therefore implicitly imposes uncon- trolled regularization on the original transformer model. On the other hand, the regularisation introduced in our technique is motivated by averaging consecutive integra- tion approximations of an ODE in the original transformer, which leads to consistent performance gain across different datasets.
is clear from Fig. 3 that the BDIA training technique leads to the best validation performance in the end of training while direct fine-tuning of GPT2 exhibits the over-fitting issue. It is plausible that the validation losses for LoRA are larger than those for FT-BDIA because LoRA only trains a small number of parameters.
It is also clear from Table 2 that Rev is most memory- efficient. BDIA-ViT improves on the validation perfor-
Table summarizes results for 5 metrics. It is seen from
7

---

On Exact Bit-level Reversible Transformers Without Changing Architecture
Table 2. Validation accuracy (in percentage), peak memory consumption, and average training time per epoch for training three models
over CIFARIO and CIFARI00. {hh ( K-1 in the BDIA training procedure were drawn from { { per training sample. The peak memory includes both the model parameters and the training states for batchsize of 128. The dropout rate was set to 0.1. See also Fig. 7
for training curves on CIFAR-IO 10 plotted against wall-clock time.
RevVIT (Mangalam et al., 2023)
VIT
BDIA-VI VIT
val. acc.
peak memory
tainn time per epoch
val. ncc.
peak memory
train. time per epoch
val. acc.
peak memory
train. time per epoch
CIFARIO
86,22 0.42
572.7MB
4
08.15:10.55
1570,6MB
30.6
89.10 1..388
493
4
CIFARIOU
61,89-00.31
572.7MB
41.8
0..46
1570.6MB
30.6
66.09-10.880
693.4MB
47.0
for CIFAR10
for CIFAR100
0.2
-VTT
VIT
-ITT RReVil
VIT
0.18
RevVil -DIA-VVT
RevVit
0.18
RevVil
BDIA-VIT
-BDIA-VIT
BDIA-VIT
0.16
0.16
64
0.14
0.14
63
taining loss
0.12
taining loss
0.12
Val. abc
val. acc
62
0.1
86
0.1
61
0.08
0.08
60
0.06
0.06
59
0.04
0.04
0.02
0.02
100
200 epochs
300
400
100
200
300
400
100
200
300
400
100
200 epochs
300
400
epochs
epochs
Figure 4. Performance comparison of ViT, RevViT (Mangalam et al., 2023), and BDIA-Vi ViT for image classification over CIFAR10 10 and 1-1 in the training procedure of BDIA-VIT VIT were drawn from (+0.5) per training sample. The dropout rate was 0.1.
CIFARIO0. {x
mance of Viv at the cost of slightly more memory, needed to store the lightweight side information {8x jk=o of the first 4 transformer blocks. We can also conclude from the table that online back-propagation does indeed significantly reduce the training memory.
operations were turned off in BDIA-transformer. We em- phasize that the non-zero "/ values were only utilized in the training process. At the inference stage, E(7x] = 0 was used for computing the validation accuracy, which in fact reduces to the standard transformer architecture.
Lastly, we briefly discuss the average training time per epoch. It is clear from the table that both RiV and BDIA- ViT require additional training time for recomputation of the intermediate activation values in online- paliinebbackropagaationn BDIA-ViT needs slightly more time than VeV because the method needs to handle the side information, the binary random variables, and quantization. See Fig. 7 for training curves on CIFAR-10 plotted against wall-clock time.
Table 3. Impact of the {7x h parameter in BDIA-VIT (w.o. quantization and w.o. online back propagation) the validation accuracy in percentage over CIFARIO. The dropout rate was 0.1.
{7e 1
0.0
(.0.
(+0.5)
(+0.6)
CIFARIO
18.15
88. $8.79:+-0.99
89.12-0.22
88.89=0.115
Table 2 summarizes the obtained validation accuracy for different setups of the (*x K- parameters. It is clear that when {hwl > 0yK-1 the performance of BDIA- transformer improves considerably in comparison to that of the conventional ViT (i.e., corresponding to BDIA-ViT with {7h = (.0)<=t). The setup of (7h = K0 K-1 performs the best. In general, the larger the magnitude of the } K-1 parameters in [0, 0.6], the slower the training speed (see Fig. 4). In practice, one can tune the magnitude of the (7x K { parameters within J0, 0.6) to trade-off between
Ablation study regarding the fyu K-1 parameters: As we mentioned in Remark 4.1, if reversibility is not of primary concern, different values for { 1K-1 can be employed for model regularisation. We con- ducted ablation study by evaluating the impact of ?h e (0.0,+0.255 +0.5, +0.6). k = 1,..., K 1, on the valida- tion performance of doing so, both the quantization and online back-propagation
for CIFAR10. For
8

---

On Exact Bit-level Reversible Transformers Without Changing Architecture
the training speed and the validation performance.
1.6
1.8
-aansommer -DDV transomme
Hansfommr
1.4
Ablation study regarding the quantization level: We have also investigated the impact of the quantization levels of {5, 0, } } on the performance of both ViT and BDIA-ViT over CIFARIO, 10, where 1 = -2 corresponds to a very coarse quantization operation. The [= variables in BDIA- ViT were drawn randomly from {=05 per training sample.
1.5
Val. loos
1.4
1
1.3
0.8
1.2
50
100
150
200
50
100
150
200
epochs
opochs
Table 4. Impact of the quantization level in VIT and BDIA-VIT on the validation accuracy in percentage over CIFARIO.
Figure 5. Performance comparison for English to French transla- tion. {ox 1x in the BDIA training procedure were randomly drawn from (+0.5) per training sample.
quantization level
-2
ViT
87.90
87.76
84.31
BDIA-VIT
89.50
88,44
84.16
existing open-source repository (i.e., the third github link of Table 7) in our experiment. The dataset being used is from Kaggle (Kelly, 2020). The tested BIDAA-trasfformer has six transformer blocks in both the encoder and decoder, respec- tively. The BDIA update expressions were implemented in both the encoder and decoder. The performance of the conventional transformer was tested as reference.
As shown in Table 4, when the quantization level I drops from 5 to -2, the performance of ViT and BDIA-VE ViT de- creases as expected. One can also observe from the table that when / E {5, 0}. BDIA- ViT performs consistently and significantly better than VIT. When ( -2, there is a large performance drop in both methods due to the very coarse quantization effect. For that case, ViT performs slightly better than BDIA-ViT. This likely is because the very coarse quantization dominates the performance and negatively af- fects the EDIA-training technique.
Fig. 5 visualizes the training and validation losses of the two tested models. is clear that the obtained curves in Fig. 5 exhibit similar properties to those in Fig. 4. That is, even though the training loss of -rans-rranfoommer is higher than the conventional transformer across all the epochs, its validation loss is significantly lower after epoch 200.
Ablation study regarding the dropout rates: As we dis- cussed earlier, the dropout technique trains an ensemble of subnetworks, while BDIA trains an ensemble of ODE solvers. It is of great interest to investigate the joint impact of dropout and BDIA. We conducted additional experiments by setting the dropout rate to 0.0 and 0.2 when testing BDIA- ViT. Table 5 shows the obtained validation accuracies, where the results for dropout rate of 0.1 are from Table 2.
6. Conclusions
We have proposed the BDIA training algorithm for trans- formers. 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 regular- izer via set of random variables [_ K-1 one variable per transformer block. The training process can then be understood as training an ensemble of ODE solvers parame- terized by {7x k==1 K-1 Exact reversibility for lossless online paak-propagation is achieved by activation quantization and storing only lightweight side information. In the inference procedure, the variables 7A replaced with their expecta- tion T[(x} = 0, which reduces the update expression to that of a conventional transformer up to activation quantization.
Table Impact of dropout rates in ViT and BDIA-ViT on the validation accuracy percentage. In BDIA training {h f==1 were randomly drawn from {+0.5) per training sample.
R-1
dmpout rate
VIT
BDIA
VIT
BDIA
VIT
BDIA
CIFARIO
86.27
89.20
85.15
89.10
87.24
88.22
CIFARIOD
59.13
64.45
61.86
66.09
61.68
64.24
It is seen from Table 5 that for different dropout rates and different datasets, the BDIA training technique improves the validation performance consistently and considerably. In brief, the dropout and BDIA are complimentary regular- ization techniques when training transformers.
Experiments on natural language generation, translation, and image classification show that BDIAA-rrnnformer pro- duces significantly better validation performance than the corresponding baseline transformers while, at the same time, reducing training memory by performing online back- propagation. In comparison, the experiments with RevVit demonstrate that it yields either inferior or comparable val- idation performance to that of ViT. BDIA AitA--taanffmmmrr is also attractive for fine-tuning since it requires minimal ar- chitecture changes, needs very little extra memory, and is complementary to LoRA.
5.3. On training BDIA-transformer for English-French translation
Language translation is a classical natural language pro- cessing (NLP) task where the transformer shows a large performance gain over other DNN models. We adopted an

---

On Exact Bit-level Reversible Transformers Without Changing Architecture
Impact Statement
Finzi, M., Izmailov, P., Maddox, W., Kirichenko, P., and Wilson
A. G. Invertible convolutional networks. In Workshop on Invertible Neural Nets and Normalizing Flows, 2019.
Transformer architectures are widely used in various con- texts including. for example, LLMs for natural language pro- cessing and diffusion models for image and video process- ing. Fine-tuning these models with either small interme- diate data sizes for diverse downstream tasks is a common practice. Our BDIA training technique offers an attractive approach for fine-tuning. This is because BDIA trains an ensemble of ODE solvers to reduce over- wert-iiting and lower training memory requirements while simultaneously retain- ing the original transformer architecture in the inference procedure.
Gholami, A., Yao, Z., Kim, S., Hooper, C., Mahoney, M. W. and Keutzer, K. AT and Memory Wall. IEEE Micro, 44(3):33-39, 2024.
Gomez, A. N., Ren, M., Urtasun, R., and Grosse, B. The
reversible residual network: Backpropagation without storing activations. arXiv: (al.::17770458555 [cs.CV], 2017.
Grathwohl, W. Chen, S. R. T. Q. Bettencourt, J., and Duve- naud, D. Ffjord: Free-form continuous dynamics for scalable reversible generative models. arXiv:1810.01367, 2018.
Hascoct, T., Febvre, Q. W. Zhuang, Y. A., and Takiguchi, T. Layer-wise invertibility for extreme memory cost reduction cnn training. In Proceedings of the IEEE/CVF International Conference on Computer Vision Workshops, 2019.
References
Achiam, J., Adler, S., Agarwal, S., Lama Ahmad, lige Akkaya, F. L. A., Almeida, D., Altenschmidt. J.. Altman, S., Anad- kat, S., Avila, R., Babuschkin, I., Balaji, S., Balcom, V., Bal- tescu, P., Bao, H., Bavarian, M., Belgum, Bello, Berdine, J. Somaddtt--Shapro.. G., Christopher Berner, I.. B., Boiko, O.. Boyd, M., Brakman, A.-L.. Brockman, G.. Brooks, T., Brundage, M.. Button, K., Cai, T., Campbell, R., Cann. A., Carey. B.. Carlson, C., Carmichael, R., Chan, B., Chang, C., Chantzis, F., Chen, D., Chen, S., Chen, R., Chen, J., Chen, M., Chess, B., Cho, C., Chu, C., Chung H. W., Cummings, D., Cur- rier. J., Dai, Y., Decarcaux, C., Degry. T., Deutsch, N., Deville, D., Dhar, A., Dohan, D., Dowling, S., Dunning, S., Ecoffet, A., Eleti, A., Eloundou, T., Farhi, D., Fedus, L.., Felix, N., Fishman, S. P.. Forte, J. Fulford, I., Gao, L.., Georges, E., Gibson, C., Goel. V. Gogineni, T., Goh, G., Gontijo-Lpess, R., Gordon, J., Grafstein, M.. Gray. S. Greene, R., Gross, J.. Gu, S. S.. Guo, Y. Hallacy, C., Han. J.. Harris, J., He, Y., Heaton, M., Heidecke, J., Hesse, C., Hickey, A., Hickey, W., Hoeschele, P., Houghton, B., Hsu, K., Hu, S., Hu, X., Huizinga, J., Jain, S., and Jain, S. Gpt-4 technical report. arXiv:2307.m 20000098889888888 [cs.CL], 2023.
He, K., 7hang, X., Ren, S., and Sun, J. Deep Residual Learning for Image Recognition. In IEEE conference on Computer Vision and Puttern Recognition (CVPR), 2015.
Hu, J., Shen, Y., Wallis, P., Allen-Zhu, Z., Li, Y., Wang. S., Wang. L., and Chen, W. LoRA: Low-Rank Adaptation of Large Language Models. arXiv:2106.09688 [cs.CL), 2021.
Jacobsen. J.-H.. Smeulders. A.. and Oyallon. E. i-revnet deep invertible networks. In ICLR, 2018
Karras, T., Aittala, M., Alia, T., and Laine, S. Elucidating the De- sign Space of Diffusion -Based Generative Models. In 36th Con- ference on Nueral Information Processing Systems (NeurlPS). 2022.
Kelly, C. Language translation fonglish-french) 2020. URL https:/ ://w////wwww// http///ww.kggll...e//// cww/ dww.//svvv.////.
Kingma, D. P. and Dhariwal, P. Glow: Generative flow with
Behrmann, J., Grathwohl, W., Chen, R. T. Q., Duvenaud, D., and Jacobsen, J.HH. Invertible Residual Networks. In Proceedingy of the International Conference on Machine Learning (ICML), 2019.
invertible Ixl convolutions. In Advances in neural information processing systems, 2018.
Kingma, D. P.. Salimans, T., Jozefowicz, R., Chen, X., Sutskever, I. and Welling, M. Improved variational inference with inverse autoregressive flow. Advances in neural information processing systems, 29, 2016.
Chang, B., Meng, L., Haber, E., Ruthotto, L., Begert, D., and Holtham, E. Reversible Architectures for Arbitrarily Deep Residual Neural Networks, #rXiv:1709.03698 [cs.CV], 2017.
Kobyzev, I., Prince, S. J., and Brubaker, M. A. Normalizing Flows: An Introduction and Review of Current Methods. PAMI, (11): 3964-3979, 2020.
Chen, R. T. Q. Rubanova, Y., Bettencourt, J., and Duvenaud, D. Neural ordinary differential equations. In 32nd Conference on Neural Information Processing Systems (Neurips), 2018
Lin, Y., Ma, X., Chu, X., Jin, Y., Yang, Z., Wang, Y., and Mei, H. LoRA dropout a sparsity regularizer for overfitting control. arXiv:2404.096 006100 [cs.LG), 2024.
Dinh, L.., Krueger, D., and Bengio. Y. Nice: Non-linear indepen- dent components estimation. arXiv preprint arXiv: Mixi::1110 410.8516, 2014.
Ma, S. The Era of 1-bit LLMs: All Large Language Models are in 1.58 Bits. arXiv:2402.17764v 2024.
Dinh, L., Sohl- Dickstein, J., and Bengio, S. Density estimation using real nvp. arXiv preprint (60500803,,......800800060333661 2016.
Mangalam, K., adn Y. Li, H. F., Qu, C.-Y,, Xiong, B., Feicht- enhofer, C. and Malik, J. Reversible vision transformers. arXiv:2302.048699 [cs.CV], 2023.
Dupont, E., Doucet, A., and Teh, W. Augmented neural odes. arXiv:1904.01681v3, 2019.
Novikova, J., Dusek, O., and Rieser, V. The e2e dataset: New challenges for end-to-end generation. arXiv preprint 1rXiv:1706.092555 2017.
Feng, J. and Huang, D. Optimal gradient checkpoint search for arbitrary computation graphs. In CVPR, 2021.
10

---

On Exact Bit-level Reversible Transformers Without Changing Architecture
Rezende, D. and Mohamed, Variational inference with normal- izing flows. International conference on machine learning, 5530-1538, PMLR, 2015.
Zhu, T. and Mangalam, Pareprop: Fast parallelized reversible backpropagation. (tt::23006009342v [cs.LG), 2023.
Sander, M. F.,, Ablin, M. B., and Peyre, G. Momentum residual ncural networks. arXiv:2102.0787, 2021.
Song. J., Meng, C., and Ermon, S. Denoising Diffusion Implicit Models. In ICLR, 2021.
Song, Y., Meng, C., and Ermon, S. Mintnet: Building invertible neural networks with masked convolutions. arXiv:1907.077 07794,, 2019.
Srivastava, N., Hinton, G., Krizhevsky, A., Sutskever, I., and Salakhutdinov. R. Dropout: A Simple Way to Prevent Neu- ral Networks from Overfitting. Journal of Machine Learning Research, pp. 1929-1958, 2014.
Stam, J. An exact bitwise reversible integrator. arXiv:2207.0769) GR], 2022.
Touvron, H., Martin, L., K. Stone, P. A., Almahairi, A., Babaei, Y.. Bashlykov, N., Batra, S., Bhargava, P., Bhosale, S., Bikel, D.. Blecher, L., Ferrer, C. C., Chen, M., Cucurull, G., Es- iobu, D., Fernandes, J., Fu, J., Fu, W., Fuller, B., Gao, C., Goswami, V. Goyal, N., Hartshor, A., Hosseini, S., Hou, R., Inan, H., Kardas, M., Kerkez, V., Khabsa, M., Kloumann, I., Korenev, A., Koura, P. S., Lachaux, M.-A., Lavril. T., Lee, J., Diana Liskovich, Y. L., Mao, Y., Martinet, X., Mihaylov, T., Mishra, P. Molybog, I,, Yixin Nie, P. Reizenstein, Rungta, R., Saladi, K. Schelten, A., Silva, R., Smith, E. M., Subrama- nian, R., Tan, X. E., Tang, B., Taylor, R., Adina Williams, J. X. K., Xu, P. Yan, Z. Zarov, I., Zhang. Y., Fan, A., Kam- badur, M., Sharan Narang, A. R., Stojnic, R., Edunov, and Scialom, Llama 2: Open foundation and fine-tuned chat models. ar arXiv:2307.0990 .09288 [cs.CL], 2023.
Wallace, Gokul, A., Ermon, S., and Naik, N. End-to-Fnd Diffusion Latent Optimization Improves Classifier Guidance. urXiv:2303.13703: [cs.CV], 2023.
Wang, N., Choi, Brand, D., Chen, C.-Y., and Gopalakrishnan, K. Training deep neural networks with 8-bit floating point numbers. In NeurlPS, 2018.
Wu, K., Park, J. B., Zhang, X., Hidayetoglu, M., Mailthody, V. S., Huang. S., Lumetta, S., and Hwu, Ssdtrain: An activation offloading framework to ssds for faster large language model training. arXiv:2408. 10013 [cs. DC), 2024.
Wu, S., Li. G., Chen, F., Shi, L. Training and inference with integers in deep neural networks. In /CLR, 2018. URL https: :/ //ppnreeviiw net / forum? forum?id=:JGxxmmppp
Yang, G., Zhang, T., Kirichenko, P., Bai, J., Wilson, G., and Sa,
C. SWALP: Stochastic Weight Averaging in oo-Prrccision Training. In ICML, 2019.
Zhang, G. On Suppressing Range Adaptive Stepsizes of Adam to Improve Generalisation Performance. ECMI, 2024.
Zhang, G.. Lewis, J. P., and Kleijn, W. Exact diffusion inversion via bidirectional integration approximation. arXiv:2307. 10829 [cs.CV], 2023a.
Zhang, G., Niwa, K., and Kleijn, W. B. On Accelerating Diffusion- Based Sampling Processes by Improved Integration Approxi- mation. arXiv:2304. 11328 [es.LG]] 2023h.

---

On Exact Bit-level Reversible Transformers Without Changing Architecture
A. Demonstration of Reversibility Issue of BDIA Without Quantization on Log Scale
note that the reconstruction errors in Fig. are plotted on a linear-scale. better visualize the results, we include Fig. for the case BDIA without quantization by following (16). The errors are plotted on log scale.
109
rocon. orror
10
10
15
270
221
222
303
22
37
277
3010
blockwise activations
Figure 6. Demonstration of the accumulated reconstruction error in log- log-scale by following (16) with the setup 74 -0.5).
k 1,...,KK K-1, when training BDIA-GPT2 with 12 transformer blocks.
B. Proof of Proposition 4.3
Proof. Equation (23) in Proposition 4.3 can be proved by simple algebra. is known from Lemma 4.2 that 2 [2k-]] = x holds for any k 2 For each element of the expression on the left hand side (LHS) of (23), we have
Q1 + sk-1[m]]
= roound Yh ( [m)]2 + 8 [m]] (m))2
(26)
By using (17) and the definition of 8;-1 [m]] from (20), it is immediate that the scalar (mx-1[ml/2" ( in (26)
is an even integer number. Since * E 0.5,, 0.5), the round operation in (26) has no real effect and can be removed. Therefore, (26) can be rewritten as
2117 (mm + s (m-1(m]2-1)
[2 [m]]22 + 8 (ml) 2
7h (ao- [m] + 8k-1 s [m|22
(27)
The proof is complete.
C. Proof of Proposition 4.4
Proof. Also (24) in Proposition 4.4 can be proved by simple algebra. Firstly, plugging (23) into (21) produces
0k+1 = (#x Sk-12-1) + (
(28)
1st term
2nd term
3rd temm
is known from Lemma 4.2 and (20) that all three terms in (28) have fixed-point precision of 2-1, As a result, the equality in (28) holds without any error. Multiplying with 7 122 on both sides of (28) and rearranging the quantities produces
1 ei((1 (1 7x) a +1+++(111 (1+ ha (xx));
- 1e41
S 2-
(29)
7k
7
which again holds without any error. The proof is complete.
D. Training Time Comparison
This Appendix discuses training effort more detail. Table 2 includes the average training time per epoch for CIFAR10 and CIFARIO0 for three algorithms and that Fig. shows the training and validation curves over 400 epochs. Fig. 7 in this Appendix shows the training curves for CIFAR1O against wall-clock time for 400 epochs.
12

---

On Exact Bit-level Reversible Transformers Without Changing Architecture
It is seen from Fig. 7 that BDIA-VIT with online backpropagation needs slightly more time to cover 400 epochs than Rev with online backpropagation. This is because BDIA-ViT needs to handle the lightweight side information, binary random variables, and quantization. VIT is most time efficient as it does not need to recompute the intermediate activation values as the other two methods.
While the final training loss of BDIA-ViT higher, as indicated in Fig. 7, its validation accuracy is better than those of the other two methods (see Fig. 4). The above property can explained by the fact that BDIA trains an ensemble ViT models using different ODE solvers and thus can be viewed complimentary regularization technique dropout.
0.25
ViT
RevViT with online backpropogation
0.2
BDIA-VIT with online backpropogation
taining loss
0.15
0.1
0.05
0
50
100
150
200
250
300
time (min)
Figure 7. Training time comparison of ViT, RevViT with online backpropagation (Mangalam et al., 2023), and BIIA-VIT with online backpropagation image classification over CIFARIO. Each training curve corresponds to 400 training epochs. Note that the BDIA technique can outperform in validation accuracy even when the training loss is higher, as seen in Fig. 4.
E. Evaluation of Joint Impact of LoRA and BDIA
It has become common practice to adopt the LoRA adapter when fine-tuning an LLM with either small intermediate data sizes. One major advantage of using LoRA is that it accelerates the training process with limited computational resources as it requires only relatively small number of parameters to be trained. Since the original LLM model parameters are kept frozen throughout the process, the knowledge learned in the pre-training is less likely to be forgotten. On the other hand, it has been reported in the literature that, even with the LoRA approach, overfitting can still occur (Lin et al., 2024).
We performed an additional experiment studying the joint impact of LoRA and BDIA when fine-tuning GPT2 M for NLG. Since the validation performance our primary concern the LoRA+BDIA training procedure was implemented without online back propaaaopan and quantization. Differently from Subsection 5.1, (7k) K in the LoRA+BDIA procedure were randomly drawn from (0.5, 0, 0.5) per training sample.
k=1
Next we briefly discuss the setup for LoRA's parameter in the LoRA+BDIA procedure. In order for the BDIA approach to have a significant effect on the performance, it is recommended to use large values for the parameter. This is because a large a value in LoRA resembles the full fine-tuning scenario to a certain extent. is known from the main paper that BDIA is designed to average every two consecutive integration approximations in a random manner via {7* A large parameter allows the averaging operations in BDIA to have a substantial effect. In the experiment, tested range of configurations with a {512, 1024] and LoRA rank E (32, 128).
Fig. 8 visualizes the training and validation performance for both LoRA and LoRA+ Similarly to the results in the main paper, the training loss of LoRA+BDIA is higher than that of LoRA alone due to the regularization effect of BDIA. On the other hand, the validation loss of LoRA+BDIA is lower than that of LoRA at the end of training, which is consistent with our expectation.
Table 6 summarizes the quality of the generated texts for four training setups evaluated with five different metrics. clear
13

---

On Exact Bit-level Reversible Transformers Without Changing Architecture
rank = 32
a 1024
LORA LORA+BDIA
1.25
LORA LORA+BDIA
2.6
train loss
val loss
1.2
2.4
1.15
2.2
1.1
0.5
1.5
2
2.5
0.5
1.5
2
2.5
iterations
x10
iterations
x x10
rank 128
a = 512
1.25
2.7
LORA LORA+BDIA
LORA LORA+BDIA
train loss
2.6
val loss
1.2
2.5
2.4
1.15
2.3
1.1
0.5
1
1.5
2
2.5
0.5
1.5
2
2.5
iterations
x 10
iterations
X10
Figure 8. Performance comparison of LoRA nad LoRA+BDIA when fine-tuning GPT2 Medium for the E2E challenge. BDIA training procedure were randomly drawn from {0.5, -0.5) per training sample.
{74)
KK11 km1 in the
Table Performance LoRA and LORA+BDIA for fine-tuning GPT-2 medium (M) E2E NLG Challenge. For all metrics, higher is better. (yx)k-7 in the BDIA training procedure were randomly drawn from (0.5, 0, -0.5) training sample.
Method
BLEU
MЕT
ROUGE-L
CIDEr
rank 32
1024
LORA+BDIA
LoRA
68.0
8.58
46.1
70.9
2.49
512
LORA+BDIA
69.0
8.70
46.5
71.5
2.52
that for all tested scenarios, the LORA+BDIA training procedure outperforms the LoRA procedure for almost all metrics. This suggests that the BDIA training approach can help LoRA in fine-tuning transformer-basee LLMss.
F. On Training BDIA-GPT2
In this experiment, we train BDIA-GPT on the openwebtext dataset utilizing the 4th github link of Table 7. 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 very small training dataset. doing so, we only took a small (i.e., 0.05%) subset from the entire dataset
14

---

On Exact Bit-level Reversible Transformers Without Changing Architecture
when training the model. The performance of GPT2 evaluated reference. Both models have 12 transformer blocks. We did not perform hyperparameter exploration and it is possible that the performance of BDIA-GPT2 would improve with such exploration.
Fig. 9 shows the training and validation curves for the two models. Similarly to Fig.4-5, BDIA-GPT2 exhibits 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 training. However, at the end of training, the validation loss of BDIA-GPT2 significantly lower than that of GPT2. This demonstrates that BDIA indeed alleviates the over-fitting issue of GPT2 for a very small training dataset.
---BII
taiiing lose
V8l. loss
8
2-
0
10000
10000
fterations
forations
Figure 9. Performance comparison when training GPT2, f7x ) (+0.5) per training sample.
in the training procedure of BDIA-ViT randomly drawn from
G. Repositories in the Experiments
Table 7, Repositories being used in the experiments
image classification
natural language generation (NLG)
language translation
nanoGPT
15


Reviews:
[{"review_id": "Pio4a9K1a8", "reviewer": "Reviewer_BqRX", "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.\n\n----\nUpdate after the rebuttal:\n----\n\nI have carefully read the authors' rebuttal as well as the comments from the other reviewers. After consideration, I have decided to increase the score. I hope the authors integrate the suggestions from the reviewers into the next version of the paper and also release the code.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.  \n\n__(1)  Regarding the paper readability__\n\nBased on  comments from all four reviewers, we will make the following changes in the revision to improve readability: \n\n(a)  We will shorten the abstract to make it more concise. \n\n(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.  \n\n(c) The pseudocode for the BDIA training method will be included in the paper to enable the readers to easily understand the method. \n\n(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.  \n\n__(2)  Regarding the reproducibility and pseudocodes of BDIA__\n\nIn 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. \n\n__(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.__\n\n\nIt 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.\n\n__(4) Regarding the small datasets in the experiments__\n\nMany 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.     \n\nSecondly, 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.\n\n__(5) Regarding impact of dropout__\n\nFirstly, 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. \n\nSecondly, 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.\n\n|dropout | 0.0 |0.0 |\n|--------------|---------------|----------| \n| |ViT|BDIA-ViT|  \n| CIFAR10| 86.27|__89.20__|\n|CIFAR100| 59.13|__64.45__|\n|dropout |0.2 | 0.2 |\n| |ViT|BDIA-ViT|  \n| CIFAR10| 87.24|__88.22__|\n|CIFAR100|61.68|__64.24__|\n\n__(6) Regarding figure 2__\n\nWe 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.  \n\n__(7) Regarding figure 1__\n\nWe will revise the caption of figure 1 to better explain how it is generated.", "claims_and_evidence": "There is a problem. Although the paper provides detailed mathematical derivations, they appear to be more closely aligned with concepts from residual networks (ResNets) rather than being specifically tailored to transformers.", "methods_and_evaluation_criteria": "Yes", "theoretical_claims": "Yes", "experimental_designs_or_analyses": "Yes. I checked all the experiments.", "relation_to_broader_scientific_literature": "It is a relevant contribution to the literature as the investigated problem is timely: reducing the memory consumption during the training of transformers.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "Strengths:\n- 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.\n- 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.\n\nWeaknesses:\n- The paper is difficult to follow. For example, the abstract is too long.\n- The reproducibility is low as there is no source codes or pseudo codes or detailed algorithms.\n- 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.\n- The experiments mainly consider small datasets or relies on toy examples for transformers.", "other_comments_or_suggestions": "No", "questions_for_authors": "- In figure 1, and in line 287, how did the authors integrate into standard transformers?\n- In figure 2, the authors should show the reconstruction errors w.r.t. the proposed method using quantization and side information. Otherwise, it is not clear the effectiveness of these tricks.\n- The authors should compare experimentally the proposed methods against vanilla transformers applied with dropout.\n- 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.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "KHm4o1lxBE", "reviewer": "Reviewer_7ZwY", "summary": "This paper introduces BDIA-transformer, an exact bit-level reversible transformer that maintains the standard architecture for inference while reducing memory consumption during training. The approach adopts bidirectional integration approximation (BDIA), allowing the authors to consider each transformer block as an exactly reversible flow (once corrected by a 1-bit stored value at each activation layer) -- a framing that is theoretically and practically important in the study of ordinary differential equations. This reversibility provides a mechanism to recompute activations during the backwards pass when training the DNNs.\n\nMore specifically, the key contributions of the paper are: (1) introducing a random hyperparameter $\\gamma \\in [-0.5,0.5]$ per transformer block per training sample, which regularizes the model by effectively training an ensemble of ODE solvers; and (2) implementing activation quantization to enable exact bit-level reversibility. During inference, $\\gamma$ is set to its expectation (zero), which reduces the architecture to a standard transformer with only activation quantization applied.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.  \n\n__(1) Regarding the computational overhead introduced by BDIA-transformer.__\n\nIn 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. \n\n__(2) Regarding a discussion of RevViT having better memory efficiency than BDIA-transformer__\n\nWe will elaborate in the paragraph at lines 369-376 in the experimental section regarding the benefit of RevViT with respect to memory efficiency. \n\n__(3) Regarding the impact of different quantization levels__\n\nAs 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.  \n\nAs 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.    \n\n |quantization level |5|0|-2|\n| ------------------------ | ----- | --- | --- |\n |ViT|87.90|87.76|__84.31__| \n|BDIA-ViT | __89.50__ |__88.47__ |84.16|  \n\n\n__(4) Regarding integrating BDIA-transformers with other memory-saving techniques__\n\nThere 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. \n\nIn 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.\n\n\n__(5)  What specific challenges did you encounter when implementing BDIA for different transformer architectures?__\n\nMany 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.  \n\nFrom 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.    \n\n\n__(6) Regarding comparisons against other memory-saving techniques beyond RevViT.__ \n\nMany 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\n\n$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.$\n\nIt 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.", "claims_and_evidence": "The authors claim their approach enables exact bit-level reversibility. This is supported by the mathematical formulation, though direct empirical verification of this property is somewhat limited in the experimental section.\n\nThe paper claims the core inference architecture is preserved by their technique. This is adequately supported by equations demonstrating how $E(\\gamma)=0$ results in a standard transformer update with only quantization added.\n\nThe paper claims BDIA-transformer reduces memory consumption during training. This is supported by memory measurements shown in Table 1, which demonstrates reduced peak memory usage compared to standard transformers.\n\nThe authors claim BDIA-transformer outperforms conventional transformers due to regularization effects. The evidence in Figures 3-6 and Tables 1-3 supports this, showing improved validation performance across different tasks.", "methods_and_evaluation_criteria": "The methods and evaluation criteria (application domains, baseline architectures, and $\\gamma$ ablation study) are generally sufficient.  Given the computational claims of the paper, metrics demonstrating such overheads would make for a more comprehensive analysis.", "theoretical_claims": "Yes.  The claims are derivational in nature and are reasonable.", "experimental_designs_or_analyses": "It would be helpful to see comparisons against other memory-saving techniques beyond RevViT.", "relation_to_broader_scientific_literature": "The authors provide a good overview of reversible neural networks, and ODE-based neural networks.  The Bidirectional Integration Approximation is reviewed.  Well known quantization strategies are mentioned.", "essential_references_not_discussed": "None.", "other_strengths_and_weaknesses": "The novelty and demonstrated effectiveness of the approach make this a strong paper.\n\nOne weakness is the lack of detailed computational overhead analysis.", "other_comments_or_suggestions": "The paper would benefit from a more thorough analysis of the computational overhead introduced by BDIA-transformer. While memory savings are quantified, training time comparisons would provide a more complete picture.\n\nThe paper mentions that RevViT has better memory efficiency than BDIA-transformer. A more nuanced discussion of this trade-off would strengthen the paper.\n\nThe ablation studies could be expanded to include analysis of the impact of different quantization levels (beyond the 9-bit quantization considered).", "questions_for_authors": "Could you elaborate on how BDIA-transformers might be integrated with other memory-saving techniques like gradient checkpointing, selective activation recomputation, and activation offloading? Are there specific combinations you've explored or would recommend?\n\nHow sensitive is the method to different quantization levels, and what are the performance trade-offs?\n\nWhat specific challenges did you encounter when implementing BDIA for different transformer architectures, and how might these insights guide others implementing your approach?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "PUl2QBYpPk", "reviewer": "Reviewer_8yrY", "summary": "The paper introduces the BDIA-transformer, a novel reversible transformer that maintains the standard transformer architecture during inference while leveraging a technique called bidirectional integration approximation (BDIA) for reversibility. The key idea is to treat each transformer block as an Euler integration approximation for solving an ordinary differential equation (ODE) and then apply BDIA to achieve reversibility. Additionally, activation quantization is incorporated to ensure exact bit-level reversibility. Experimental results on image classification, natural language generation, and language translation tasks show that BDIA-transformers outperform conventional transformers in validation accuracy while significantly reducing training memory.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.  \n\n__(1) Regarding the impact of different quantization levels on performance.__\n\nWe 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.  \n\nAs 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.   \n\n |quantization level |5|0|-2|\n| ------------------------ | ----- | --- | --- |\n |ViT|87.90|87.76|__84.31__| \n|BDIA-ViT | __89.50__ |__88.47__ |84.16|  \n\n__(2) Regarding a discussion on the trade-off between memory savings and computational overhead__\n\nIn 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. \n\n__(3) Regarding the impact of quantization level on memory savings__\n\nFirstly, 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.    \n\nSecondly, 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.    \n\n__(4) Can BDIA be applied to other transformer variants?__\n\nMany 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. \n\nWe 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.\n\n||CIFAR10|  \n|---------------------|------|\n|vision-linformer|84.16|   \n|BDIA-vision-linformer|__86.6__|  \n\n__(5)  What are the potential limitations of BDIA when scaling to large models?__\n\nThese 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. \n\nOne 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. \n\n__(6) Regarding literature review on linear attention mechanisms and memory-efficient self-attention__\n\nWe will update the introduction accordingly to reflect the works you mentioned.", "claims_and_evidence": "The paper claims that BDIA-transformers achieve exact bit-level reversibility without altering the standard transformer architecture during inference. The proposed method reduces memory consumption in training while improving validation accuracy through model regularization. Experimental evidence supports these claims, showing improved validation performance across multiple tasks and datasets. The memory savings are attributed to online back-propagation enabled by reversibility, and the performance gain is linked to the ensemble effect of ODE solvers induced by the stochastic nature of the hyper-parameter. The claim that BDIA-transformers outperform RevViT is also supported by empirical results, though the performance gain varies across datasets.", "methods_and_evaluation_criteria": "The evaluation is conducted on tasks including image classification, natural language generation, and language translation. The metrics used include validation accuracy, training loss, and peak memory consumption. The use of CIFAR-10, CIFAR-100, and the E2E dataset provides diverse benchmarks to assess the generalization of the proposed method. The evaluation criteria are appropriate for demonstrating the advantages of memory efficiency and performance improvements in BDIA-transformers compared to standard transformers and RevViT.", "theoretical_claims": "The theoretical claim  is supported by the mathematical formulation of the update expressions. The derivation of BDIA update rules and the use of activation quantization ensure that each forward and backward step can be losslessly reconstructed. The use of lightweight side information per transformer block to counteract binary quantization loss further strengthens the theoretical soundness.", "experimental_designs_or_analyses": "Experimental results compare the proposed model against baseline transformers and RevViT across various tasks. Multiple datasets and evaluation metrics are adopted. The training setup, involving multiple experimental repetitions, mitigates the effect of randomness. The inclusion of ablation studies exploring the impact of different hyper-parameters and quantization settings provides additional insights into the model's behavior.", "relation_to_broader_scientific_literature": "The paper builds on prior work in reversible neural networks and ODE-based modeling, extending these ideas to transformers. References to RevNet, neural ODEs, and RevViT contextualize the contributions within the broader literature. The connection to diffusion models and integration approximations further enriches the theoretical foundation.", "essential_references_not_discussed": "A deeper discussion on recent advances in memory-efficient transformer training techniques would strengthen the paper quality. Notably, works on linear attention mechanisms and memory-efficient self-attention could provide alternative perspectives on reducing memory consumption.", "other_strengths_and_weaknesses": "Strengths:\n\n- Novel use of BDIA in transformers, achieving exact bit-level reversibility without modifying the inference architecture.\n\n- Significant memory savings and improved validation performance.\n\n- Thorough experimental evaluation across diverse tasks. Insightful ablation studies on the hyper-parameter and activation quantization.\n\nWeaknesses:\n\n- Increased memory consumption due to storing lightweight side information.\n\n- Dependence on quantization introduces complexity in implementation.", "other_comments_or_suggestions": "- Can you clarify the impact of different quantization levels on performance?\n\n- A discussion on the trade-off between memory savings and computational overhead would be helpful.", "questions_for_authors": "1. How does the choice of quantization level affect the model's performance and memory savings?\n\n2. Can BDIA be applied to other transformer variants, such as linear transformers or long-form transformers?\n\n3. What are the potential limitations of BDIA when scaling to large models with more than 100B parameters?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BafWTF2FAq", "reviewer": "Reviewer_tmy4", "summary": "The paper proposes BDIA-transformer, a novel approach combining the bidirectional integration approximation (BDIA) method and activation quantization to achieve exact bit-level reversibility in standard transformer architectures. This combination significantly reduces memory usage during training via online back-propagation, where activations are recomputed on-the-fly rather than stored. Additionally, random sampling of the hyperparameter $\\gamma  \\in ${-0.5, 0.5} introduces regularization that empirically improves generalization. Experiments on ViT-based image classification (CIFAR-10/100), GPT2-based language translation show BDIA-transformer achieving better validation performance and reduced training memory compared to baseline transformers.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.  \n\n__(1) Clarify the theoretical underpinnings of the reversibility claim explicitly.__\n\nIn 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, \nwhich 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\\\\}$.\n\n__(2) Include sensitivity analysis regarding quantization precision__\n\nAs 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.  \n\nAs 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.  \n\n |quantization level |5|0|-2|\n| ------------------------ | ----- | --- | --- |\n |ViT|87.90|87.76|__84.31__| \n|BDIA-ViT | __89.50__ |__88.47__ |84.16|  \n\n__(3) Regarding the relationship between  quantization precision and reversibility.__\n\nIn 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.   \n\n\n__(4) Consider additional experiments explicitly varying model size and data to highlight potential trade-offs__\n\nMany 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.    \n\nWe 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.   \n\n\n\n| number of transformer blocks | 6|10|   \n|-----------------------------------|----|----|\n|ViT|88.15 |88.02|  \n|BDIA-ViT|__89.10__|__89.08__|  \n\n\n\n__(5) Regarding the applicability of BDIA to other DNN blocks__\n\nReviewer 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.\nWe 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. \n\n\nDue 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.\n\n||CIFAR10|  \n|---------------------|------|\n|vision-linformer|84.16|   \n|BDIA-vision-linformer|__86.6__|", "claims_and_evidence": "Most claims concerning improved validation accuracy, reduced memory consumption, and generalization benefits are supported by experiments. While implicitly verified, it would be great to update Fig 2. to add the reconstruction errors when using a model following Equation 24 in order to confirm the validity of the exact byte inversion.", "methods_and_evaluation_criteria": "The methods and evaluation criteria generally make sense for the demonstrated applications. CIFAR10/100, GPT2 fine-tuning, and English-French translation tasks are adequate benchmarks to show memory efficiency and regularization effects. However, these memory considerations often arise in large-scale settings where the model's capacity is the limiting factor, but the experiments are done in a context where data is the limiting factor (overfitting setting).", "theoretical_claims": "No formal proofs were provided to substantiate theoretical claims, particularly regarding exact reversibility under quantization, leaving this as a significant gap.", "experimental_designs_or_analyses": "The experimental designs appear valid, clearly illustrating the benefits of BDIA-transformer under limited data conditions. Nonetheless, an analysis of scenarios where model capacity, rather than data quantity, limits performance would help clarify the applicability and limitations of the proposed regularization method. While large-scale experiments are not expected, toy experiments could give a good picture of how the method trades model capacity to improve generalization.", "relation_to_broader_scientific_literature": "The paper positions itself appropriately within literature on reversible neural networks and activation quantization, clearly differentiating its unique combination of these two ideas from existing approaches like RevNet and RevViT.", "essential_references_not_discussed": "To my knowledge, the paper adequately cites and discusses relevant existing literature;", "other_strengths_and_weaknesses": "Strengths:\n- Novel combination of reversible architectures with quantization for memory efficiency and regularization.\n- Empirical validation across multiple popular tasks demonstrates general applicability.\n- The fact that it minimally changes the architecture allows the use of this method for finetuning, which is a creative and powerful use case.\n\nWeaknesses:\n- Missing formal proofs regarding exact reversibility under quantization.\n- No detailed exploration of how quantization precision (parameter l) impacts performance and memory. It would help to disentangle the benefits from reversibility and the benefits from quantization.\n- Limited experimental settings that explicitly test the potential for underfitting or performance degradation when model capacity is constrained.", "other_comments_or_suggestions": "- Clarify the theoretical underpinnings of the reversibility claim explicitly.\n- Include sensitivity analysis regarding quantization precision.\n- Consider additional experiments explicitly varying model size and data to highlight potential trade-offs.", "questions_for_authors": "- I am curious about how sensitive your method is to the choice of quantization precision (l). Especially how this precision affects both performance and reversibility?\n- I don't see particular counter-indications to applying this strategy to blocks other than attention. Can similar observations be made for convolutional networks?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This manuscript targets the problem of reducing memory use during training of (large scale) transformer models. This is accomplished through the use of a bidirectional integration approximation (BDIA). This is complemented by quantization.

The claims of the manuscript are well supported by the experiments (though the scale is, somewhat necessarily, limited). The reviewers like the underlying ideas behind the work and there does seem to be real potential for impact. Nevertheless, reviewers did identify that the presentation could be improved and more concrete theoretical claims would strengthen the manuscript (which are supposedly forthcoming as minor revisions). The interesting and potentially impactful core ideas in the manuscript paired with its good empirical evaluation justify my recommendation.
Decision: Accept (poster)