Dr$^2$Net: Dynamic Reversible Dual-Residual Networks for Memory-Efficient Finetuning

TL;DR

The paper tackles the memory bottleneck of end-to-end finetuning large pretrained vision models on high-dimensional data. It introduces Dynamic Reversible Dual-Residual Networks (Dr$^2$Net), a surrogate reversible backbone that preserves the pretrained residual structure while adding a reversible residual path controlled by coefficients $oldsymbol{\alpha}$ and $oldsymbol{eta}$. A dynamic finetuning strategy gradually morphs from the original non-reversible model to a robust reversible network, balancing initialization fidelity and gradient precision to maintain performance. Empirical results across multiple tasks show substantial memory savings with accuracy comparable to conventional finetuning, validating the approach for memory-constrained settings and high-resolution data scenarios. Overall, Dr$^2$Net offers a practical, scalable pathway to memory-efficient finetuning with potential applicability beyond vision to other domains.

Abstract

Large pretrained models are increasingly crucial in modern computer vision tasks. These models are typically used in downstream tasks by end-to-end finetuning, which is highly memory-intensive for tasks with high-resolution data, e.g., video understanding, small object detection, and point cloud analysis. In this paper, we propose Dynamic Reversible Dual-Residual Networks, or Dr$^2$Net, a novel family of network architectures that acts as a surrogate network to finetune a pretrained model with substantially reduced memory consumption. Dr$^2$Net contains two types of residual connections, one maintaining the residual structure in the pretrained models, and the other making the network reversible. Due to its reversibility, intermediate activations, which can be reconstructed from output, are cleared from memory during training. We use two coefficients on either type of residual connections respectively, and introduce a dynamic training strategy that seamlessly transitions the pretrained model to a reversible network with much higher numerical precision. We evaluate Dr$^2$Net on various pretrained models and various tasks, and show that it can reach comparable performance to conventional finetuning but with significantly less memory usage.

Figures (6)

• Figure 1: Comparison of different ways of finetuning from pretrained non-reversible models. (a) Conventional finetuning uses the same non-reversible architecture in the downstream task, initialized with the pretrained parameters. It consumes high GPU memory. (b) Previous reversible methods (e.g., revnetrevvitzhao2023re2tal) cannot finetune from pretrained non-reversible models on the downstream task due to architecture discrepancy. They show reduced accuracy when training from scratch on the downstream. (c) Our proposed Dr$^2$Net can directly finetune from pretrained non-reversible networks, significantly saving memory while preserving accuracy. The top-right chart illustrates memory usage and accuracy for temporal action detection on ActivityNet-v1.3 caba2015activitynet using VSGN zhao2021video and Video Swin vswin.
• Figure 2: Transforming a pretrained non-reversible network architecture $\mathcal{M}_n$ into our proposed Dr$^2$Net.(a) $\mathcal{M}_n$: the pretrained non-reversible network with three blocks $\mathcal{F}_i, i=1,2,3$. Considering that most contemporary networks have residual connections, we illustrate the network with residual connections in the figure (green arrows), though our method doesn't restrict $\mathcal{M}_n$ to be residual networks. (b) DrNet: a reversible network obtained by adding a new group of residual connections (pink arrows) to $\mathcal{M}_n$. (c) Dr$^2$Net: our proposed reversible network obtained by adding coefficients $\alpha$ and $\beta$ to the two groups of residual connections respectively. Dr$^2$Net is equivalent to $\mathcal{M}_n$ when $\alpha=1$ and $\beta=0$. Note that the blocks $\mathcal{F}_i$ can be of any architectures following zhao2023re2tal, and there can be any number of $\mathcal{F}_i$ blocks in each network.
• Figure 3: Gradient error levels with different $\alpha$ and $\beta$ values. The scales on the right colorbar represent $10^{-12} \sim 10^{-2}$. When $\alpha=1$ and $\beta=0.1$ (top-right) at the beginning of finetuning, the error level is $10^{-5}$. The errors are the smallest when $\alpha=0$ and $\beta=1$ (bottom-left) and the largest when $\alpha=1$ and $\beta=1$ (bottom-right). The error level in the middle area $10^{-7}$ is already acceptable. The blue arrows represent an ideal evolution path of the two coefficients over the finetuning process: progressively approaching the values that produce acceptable gradient errors.
• Figure 4: $\mathcal{F}_i$ blocks in a transformer network. If the pretrained model is a transformer network, e.g., Swin liu2021swin or ViT dosovitskiy2021image, the $\mathcal{F}_i$ blocks in our Dr$^2$Net are attention layers or MLP layers. The two types of layers are interleaved, namely, if $\mathcal{F}_1$ is an attention layer, then $\mathcal{F}_2$ is an MLP layer, and $\mathcal{F}_3$ is an attention layer, and so on.
• Figure 5: Forward and reverse computation in Dr$^2$Net. Gray arrows denote the pathway for $x_i$, and pink arrows denote the pathway for $y_i$. Compared to Fig. \ref{['fig:unify_res_rev']}, we place the $\mathcal{F}_i$ blocks along with their $\alpha$-weighted residual connections inside the module $\mathcal{G}_i$.