Momentum ResNets: Reversible, Efficient Deep Networks

• Momentum ResNets are deep architectures that integrate momentum dynamics to enrich trajectory behavior beyond standard ResNets.
• They achieve analytical invertibility by reconstructing previous states on demand, significantly reducing activation memory costs.
• Empirical results across CIFAR and ImageNet benchmarks demonstrate improved accuracy and efficiency while supporting complex, non-homeomorphic mappings.

Momentum ResNets are a class of deep neural architectures that augment standard Residual Networks (ResNets) with explicit momentum dynamics, yielding strictly greater representational capacity, analytical invertibility, and memory efficiency. By replacing the canonical first-order update rule of ResNets with a momentum-augmented scheme inspired by second-order ordinary differential equations (ODEs), Momentum ResNets—encompassing both m-RevNet and related formulations—integrate principles from numerical analysis, dynamical systems, and control theory. These models enable a reversible computation graph, support efficient training of very deep networks, and can represent non-homeomorphic transformations that are inaccessible to classical residual flows. The theoretical foundation and empirical performance of Momentum ResNets have been established across a range of computer vision benchmarks and are supported by rigorous analysis.

1. Mathematical Foundation and Forward Dynamics

Momentum ResNets implement layerwise updates governed by a momentum principle derived from damped second-order ODEs. The general discrete update at layer $n$ is

$$\begin{cases} v_{n+1} = \gamma v_n + (1-\gamma) f(x_n, \theta_n) \ x_{n+1} = x_n + v_{n+1} \end{cases}$$

where $x_n \in \mathbb{R}^d$ is the activation (“position”) at layer $n$, $v_n$ is the “velocity,” $\gamma \in [0, 1]$ is a momentum parameter, and $f$ is a standard residual block (e.g., a small CNN or MLP) (Sander et al., 2021). This scheme can alternatively be parameterized by a damping/momentum coefficient $\mu>0$ via

$$\begin{cases} \nu_{t+1} = (1-\mu) \nu_t + \mu \mathcal{F}(x_t, \theta_t) \ x_{t+1} = x_t + \nu_{t+1} \end{cases}$$

which arises from a semi-implicit Euler discretization of

$$\frac{1}{\mu} \frac{d^2 h(t)}{dt^2} + \frac{d h(t)}{dt} = \mathcal{F}(h(t), \theta)$$

subject to suitable initial conditions (Li et al., 2021).

The classical ResNet update $$\begin{cases} v_{n+1} = \gamma v_n + (1-\gamma) f(x_n, \theta_n) \ x_{n+1} = x_n + v_{n+1} \end{cases}$$0 is recovered in the limit $$\begin{cases} v_{n+1} = \gamma v_n + (1-\gamma) f(x_n, \theta_n) \ x_{n+1} = x_n + v_{n+1} \end{cases}$$1 or $$\begin{cases} v_{n+1} = \gamma v_n + (1-\gamma) f(x_n, \theta_n) \ x_{n+1} = x_n + v_{n+1} \end{cases}$$2. For generic $$\begin{cases} v_{n+1} = \gamma v_n + (1-\gamma) f(x_n, \theta_n) \ x_{n+1} = x_n + v_{n+1} \end{cases}$$3 and $$\begin{cases} v_{n+1} = \gamma v_n + (1-\gamma) f(x_n, \theta_n) \ x_{n+1} = x_n + v_{n+1} \end{cases}$$4, the system interpolates between first-order and second-order flow regimes, admitting richer trajectory dynamics.

2. Reversibility and Memory Complexity

A key property of Momentum ResNets is exact analytical invertibility at each block. Given $$\begin{cases} v_{n+1} = \gamma v_n + (1-\gamma) f(x_n, \theta_n) \ x_{n+1} = x_n + v_{n+1} \end{cases}$$5, the preceding state $$\begin{cases} v_{n+1} = \gamma v_n + (1-\gamma) f(x_n, \theta_n) \ x_{n+1} = x_n + v_{n+1} \end{cases}$$6 is reconstructed by

$$\begin{cases} v_{n+1} = \gamma v_n + (1-\gamma) f(x_n, \theta_n) \ x_{n+1} = x_n + v_{n+1} \end{cases}$$7

Analogous inversion holds for the $$\begin{cases} v_{n+1} = \gamma v_n + (1-\gamma) f(x_n, \theta_n) \ x_{n+1} = x_n + v_{n+1} \end{cases}$$8-parameterization (Sander et al., 2021, Li et al., 2021). In a network composed of $$\begin{cases} v_{n+1} = \gamma v_n + (1-\gamma) f(x_n, \theta_n) \ x_{n+1} = x_n + v_{n+1} \end{cases}$$9 such blocks, only the terminal state $x_n \in \mathbb{R}^d$0 needs to be stored during the forward pass. At backward time, previous states are recomputed sequentially on demand. This design ensures the total memory footprint for activations is $x_n \in \mathbb{R}^d$1 in $x_n \in \mathbb{R}^d$2, in contrast to the $x_n \in \mathbb{R}^d$3 cost for standard ResNets. For example, an m-RevNet-269 on ImageNet requires $x_n \in \mathbb{R}^d$413M activations, versus 104M for a ResNet-269 (Li et al., 2021).

Backward propagation alternates block inversion (for recomputation) and automatic differentiation of $x_n \in \mathbb{R}^d$5, with no need for checkpointing or additional memory aside from minimal “bit-loss buffers” to ensure numerical exactness in floating-point arithmetic (Sander et al., 2021).

3. Representational Power and Theoretical Properties

Momentum ResNets, as discrete second-order flows, strictly generalize first-order ODE-based ResNets. Any flow realizable by a ResNet is admissible via suitable parameterization of a Momentum ResNet, but the converse is false: Momentum ResNets can represent non-homeomorphic maps with crossing trajectories.

• Crossing Flows: First-order flows preserve topology and cannot realize mappings such as sign-flips ($x_n \in \mathbb{R}^d$6 in 1D) or nested sphere separation in higher dimensions.
• Second-order Flows: Admitting velocity enables trajectory crossing and hence richer function classes. Explicit constructions demonstrate that m-RevNets can achieve $x_n \in \mathbb{R}^d$7 via analytic solutions of suitably chosen ODEs (Li et al., 2021).
• Linear Mapping Characterization: In the continuous-time (small step) and linear setting, any linear invertible mapping up to scalar multiple is attainable by a Momentum ResNet; first-order systems are restricted to $x_n \in \mathbb{R}^d$8, forbidding many negative eigenvalue patterns (Sander et al., 2021).
• Universal Approximation: Under standard assumptions, Momentum ResNets with ReLU activation admit universal $x_n \in \mathbb{R}^d$9 approximation and simultaneous interpolation of arbitrary finite input-output matching under suitable control functions (Ruiz-Balet et al., 2021).

These properties yield empirical and theoretical superiority over vanilla ResNets on tasks where topological complexity or precise trajectory steering is required.

4. Empirical Performance and Benchmarks

Momentum ResNets have been extensively benchmarked against standard ResNets and other invertible architectures on canonical datasets:

• CIFAR-10/100: m-RevNet and Momentum ResNet achieve error rates equal to or lower than corresponding ResNet architectures, e.g., m-RevNet-110: 5.19% (CIFAR-10), 25.15% (CIFAR-100) versus ResNet-110: 5.74% / 26.44% (Li et al., 2021). On CIFAR-10, Momentum ResNet-101 matches vanilla ResNet-101 at 95.1% accuracy (Sander et al., 2021).
• ImageNet: Comparable parameter-matched m-RevNet models demonstrate lower or equal top-1 error (m-RevNet-101: 22.1%; ResNet-101: 23.0%) at a fraction of the memory cost (Li et al., 2021). Throughput and wall-clock time are competitive.
• Semantic Segmentation: On Cityscapes and ADE20K, substituting m-RevNet into PSPNet backbones substantially increases mIoU under identical or larger batch sizes (Li et al., 2021).
• Toy Tasks and Control: On pathological mappings (e.g., $n$0 in 1D, nested rings in 2D), ResNets fail due to the topological constraints of first-order flows, while Momentum ResNets succeed (Sander et al., 2021, Li et al., 2021).
• Learning to Optimize: Embedding ISTA updates in a reversible block, only the momentum-augmented version converges stably to fixed points, outperforming other invertible designs (Sander et al., 2021).

A summary of key performance metrics is provided:

Model	CIFAR-10 Error	ImageNet Top-1 Error	Activation Memory (69M params)
ResNet-269	5.24%	23.0%	~104M
m-RevNet-269	5.09%	22.1%	~13M

(Values from (Li et al., 2021); memory values are for ImageNet-sized input.)

5. Practical Implementation and Tuning

Momentum ResNets are fully compatible as drop-in replacements for traditional ResNet blocks. Standard guidelines:

• Block Replacement: Each residual block's update is replaced with the momentum-augmented rule; channel sizes and kernel dimensions are unchanged (Li et al., 2021).
• Hyperparameters: Recommended $n$1 or $n$2; default $n$3 or $n$4 yields robust performance. For initial velocity ($n$5), both zeros and data-driven initialization are possible, with marginal performance differences at the cost of extra parameters.
• Backpropagation: Use auto-differentiation (e.g., PyTorch's autograd.grad) and implement custom backward passes that re-invert each block before gradient computation (Li et al., 2021).
• Numerical Precision: Division operations in inversion require small “bit-buffer” registers to prevent round-off drift; float32 is sufficient for $n$6.
• Fine-tuning and Transfer: Pre-trained ResNets can be converted to Momentum ResNets with no architectural modification and minimal retraining, enabling memory-efficient fine-tuning on resource-constrained environments (e.g., batch size 4 versus 2 in 3GB GPU memory for medical image transfer) (Sander et al., 2021).
• Training Regime: Standard schedules (learning rate, weight decay, augmentation) are retained. Memory efficiency allows increasing batch sizes, improving batch-norm statistics and model scaling (Li et al., 2021).

6. Theoretical Analysis: Control and Approximation

The expressivity of Momentum ResNets has been formalized via control-theoretic arguments:

• Simultaneous Controllability: Arbitrary finite interpolation tasks (matching input-output pairs) are solvable by explicit piecewise-constant controls. This leverages flows in position-velocity space and the existence of contractive, oscillatory, and parallel-displacement regimes (Ruiz-Balet et al., 2021).
• Universal Approximation in $n$7: For measurable functions on bounded domains, with ReLU activation and sufficient depth, the universal $n$8 approximation property is established for the flows generated by momentum ResNets (Ruiz-Balet et al., 2021).
• Tracking Capability: Neural ODEs with auxiliary memory states (a related higher-order structure) further enable simultaneous tracking of time-dependent targets, provided the memory state is of sufficient dimension. This suggests that the memory-augmented architecture is strictly more expressive than both standard ResNets and vanilla first-order Neural ODEs.

The analysis shows that momentum (or memory) enables trajectories to collide in state space without loss of distinguishability, overcoming limitations inherent to first-order flows.

7. Outlook and Architectural Implications

Momentum ResNets unify reversible computation, memory-efficient training, and enhanced representational capacity in a single architecture, with broad implications:

• Reversible Architectures: Unlike previous invertible network designs (e.g., RevNets, i-ResNets), Momentum ResNets do not require channel partitioning, spectral constraints, or inner loops, allowing direct drop-in use in legacy codebases (Sander et al., 2021).
• Model Scaling: The memory savings facilitate training of deeper networks, larger batch sizes, or larger input resolutions without sacrificing throughput or stability (Li et al., 2021).
• Control and Dynamical Systems: The relationship to second-order ODEs and control theory provides a rigorous framework for understanding and engineering network expressivity, robustness, and trajectory design. Extensions to architectures with explicit memory (as in Memory NODEs) further increase flexibility (Ruiz-Balet et al., 2021).
• Expressive Power: The strict enrichment of function classes realizable by Momentum ResNets compared to ResNets holds for both linear and nonlinear settings and is supported by constructive proofs and empirical evaluations across multiple domains.
• Future Directions: Exploring generalizations to higher-order dynamics, fractional, or delayed flows may further enhance expressiveness and stability, with potential in both theory and applications (Ruiz-Balet et al., 2021).

The integration of momentum into residual architectures enables strict improvements in expressiveness, memory efficiency, and practical trainability, with rigorous theoretical support and demonstrated empirical benefits across standard vision benchmarks and control-inspired problems (Li et al., 2021, Sander et al., 2021, Ruiz-Balet et al., 2021).