ADMMiRNN: Efficient ADMM for RNN Training

• The paper introduces ADMMiRNN, which recasts RNN training as a block-wise constrained optimization problem via ADMM to overcome vanishing/exploding gradients.
• It leverages proximal penalties and closed-form updates to decompose the problem, resulting in robust, initialization-insensitive convergence.
• Parallel implementations—both synchronous and asynchronous—enhance scalability and speed up convergence on sequence modeling tasks.

ADMMiRNN is a training framework for Recurrent Neural Networks (RNNs) based on the Alternating Direction Method of Multipliers (ADMM). It is designed to circumvent classical gradient-based difficulties in RNN optimization, notably vanishing/exploding gradients and extreme sensitivity to initialization. ADMMiRNN reformulates RNN training as a constraint-splitting problem using the unfolded RNN computational graph and applies a block-wise ADMM approach, enabling stable and efficient convergence. Parallel extensions, including both synchronous and asynchronous variants, further address scalability and efficiency concerns. Experimental results demonstrate ADMMiRNN’s superior stability and convergence properties on standard sequence modeling tasks compared to conventional optimizers (Tang et al., 2020).

1. Motivation and Core Challenges

RNNs suffer from two central obstacles in gradient-based training: (1) vanishing gradients, in which backpropagated signals exponentially decay through the unfolded RNN as the singular values of the recurrent Jacobian fall below one, and (2) exploding gradients, where singular values above one induce exponential amplification. This dynamic is further exacerbated by sensitivity to initial weight and bias values; small variations can drastically alter the recurrent Jacobian spectrum, destabilizing training trajectories. Standard optimizers such as SGD, Adam, and RMSProp propagate gradients through each timestep, inheriting both these failure modalities.

ADMMiRNN addresses these challenges by recasting RNN training as an equality-constrained optimization problem using explicit local variables for each timestep’s pre-activation, state, and output. This enables the application of ADMM, which separates the original problem into simpler subproblems and avoids explicit backpropagation, thereby robustly mitigating vanishing/exploding gradients and improving initialization insensitivity. Proximal penalty terms in each subproblem further regularize optimization, and the structure naturally admits parallelism (Tang et al., 2020).

2. Mathematical Framework

Consider an unfolded RNN of sequence length $N$ with parameters

• $u \in \mathbb{R}^{d \times d_x}$
• $w \in \mathbb{R}^{d \times d}$
• $b \in \mathbb{R}^d$
• $v \in \mathbb{R}^{d_y \times d}$
• $c \in \mathbb{R}^{d_y}$

Defining the state evolution and output as \begin{align} a_t &= u\,x_t + w\,s_{t-1} + b \ s_t &= f(a_t) \ o_t &= v\,s_t + c \end{align} and with loss $R(o_1,\ldots,o_N)$ and regularizer $\Omega(w)$, the standard objective is

$\min_{u,w,b,v,c}\;R(o_{1:N})+\Omega(w)$

This is converted to an ADMM-suitable form by introducing explicit variables $\{a_t, s_t, o_t\}_{t=1}^N$ and constraints $u \in \mathbb{R}^{d \times d_x}$0, $u \in \mathbb{R}^{d \times d_x}$1, and $u \in \mathbb{R}^{d \times d_x}$2, which are relaxed into the augmented Lagrangian

$u \in \mathbb{R}^{d \times d_x}$3

as detailed in the definition (AL), with penalty and relaxation parameters $u \in \mathbb{R}^{d \times d_x}$4. The full variable vector is split into eight blocks: $u \in \mathbb{R}^{d \times d_x}$5, $u \in \mathbb{R}^{d \times d_x}$6, $u \in \mathbb{R}^{d \times d_x}$7, $u \in \mathbb{R}^{d \times d_x}$8, $u \in \mathbb{R}^{d \times d_x}$9, $w \in \mathbb{R}^{d \times d}$0, $w \in \mathbb{R}^{d \times d}$1, $w \in \mathbb{R}^{d \times d}$2.

Each block is updated iteratively, alternating between a backward and forward sweep, and the dual multipliers are updated according to standard ADMM rules. All subproblems in this block decomposition have strongly convex quadratic (or simple) objective functions, and there are closed-form solutions or easily computed updates for each block.

3. Block-wise Updates and Solver Structure

To maintain computational efficiency and avoid expensive matrix operations, each variable block update is linearized with a proximal term. Denote the current iteration by $w \in \mathbb{R}^{d \times d}$3, and the next by $w \in \mathbb{R}^{d \times d}$4.

The $w \in \mathbb{R}^{d \times d}$5-update exemplifies the structure: $w \in \mathbb{R}^{d \times d}$6 and is solved in closed form. The $w \in \mathbb{R}^{d \times d}$7, $w \in \mathbb{R}^{d \times d}$8, $w \in \mathbb{R}^{d \times d}$9, and $b \in \mathbb{R}^d$0 updates follow similar proximal quadratic structure. For the timestep-specific auxiliary variables $b \in \mathbb{R}^d$1, the subproblems decouple across $b \in \mathbb{R}^d$2 and can be parallelized.

The update schedule for all blocks can be summarized as follows (see Algorithm 1 in (Tang et al., 2020)):

1. Backward sweep: $b \in \mathbb{R}^d$3
2. Forward sweep: $b \in \mathbb{R}^d$4
3. Dual multiplier update

This approach removes all explicit backpropagation through time; each update step is either a small quadratic problem or an element-wise proximal computation.

4. Convergence Properties

ADMMiRNN satisfies several theoretical guarantees under standard assumptions, specifically Lipschitz continuity of the loss gradient and boundedness:

• Descent and Boundedness (Theorem 1): If each penalty $b \in \mathbb{R}^d$5, where $b \in \mathbb{R}^d$6 is the Lipschitz constant, the Lagrangian decreases at each iteration, and all iterates remain bounded.
• Global Convergence (Theorem 2): Under the same conditions, every limit point of the sequence of iterates is a stationary point of the augmented Lagrangian, satisfying KKT conditions.
• Sublinear Rate (Theorem 3): The best-so-far gap $b \in \mathbb{R}^d$7, achieving asymptotic sublinear ergodic convergence.

Asynchronous variants with bounded staleness $b \in \mathbb{R}^d$8 and strongly convex loss functions also retain $b \in \mathbb{R}^d$9 convergence to the KKT set (Theorem 4).

5. Parallel and Distributed Implementations

The independence of many subproblems at each ADMM step allows for both synchronous and asynchronous parallelization.

• Synchronous Parallel ADMMiRNN (SP-ADMMiRNN): The eight variable blocks are partitioned among $v \in \mathbb{R}^{d_y \times d}$0 workers. Updates are performed in parallel per block, with a barrier synchronization at each iteration. Workers exchange and synchronize variables at each global step.
• Asynchronous Parallel ADMMiRNN (AP-ADMMiRNN): Workers independently pull the current state, update their block, and return the result without global synchronization. Bounded staleness $v \in \mathbb{R}^{d_y \times d}$1 ensures update freshness. All communication follows a master–worker protocol, with the master handling block distribution, gathering updates, and merging iterates (see Algorithms 2 and 3 in (Tang et al., 2020)).

Table: Modes of Parallelization

Mode	Worker Coordination	Convergence Guarantee
SP-ADMMiRNN	Barrier at each sweep	Global, as in central
AP-ADMMiRNN	Asynchronous, staleness bounded by $v \in \mathbb{R}^{d_y \times d}$2	$v \in \mathbb{R}^{d_y \times d}$3 to KKT under strong convexity

6. Experimental Results

Benchmarks were conducted on MNIST (sequence-mode), IMDb sentiment analysis, and synthetic text classification tasks with variable sequence lengths. The principal findings are:

• Convergence Speed: ADMMiRNN achieves near-zero training loss within a few dozen iterations. Comparators (SGD, Adam, RMSProp) require hundreds to thousands of iterations for comparable results.
• Test Accuracy: On MNIST, over 99% accuracy is achieved within approximately 100 iterations, whereas baselines either plateau lower or converge more slowly. On IMDb, ADMMiRNN shows rapid and smooth ascent of sentiment classification accuracy.
• Stability: Across 10 independent initializations per optimizer, ADMMiRNN exhibits virtually zero standard deviation in loss/accuracy, compared to significant variability with SGD-style optimizers.
• Resilience to Sequence Length: On the synthetic text task, conventional RNNs show loss escalation with increasing sequence length, attributed to vanishing/exploding gradients. ADMMiRNN’s loss remains flat as sequence length grows.
• Parallel Efficiency: On three-worker configurations, SP-ADMMiRNN reduces runtime by approximately 40% and AP-ADMMiRNN by 70%, with convergence profiles preserved.

7. Practical Aspects and Implementation Guidelines

• Penalty Parameters: Penalties $v \in \mathbb{R}^{d_y \times d}$4 must satisfy $v \in \mathbb{R}^{d_y \times d}$5 (Lipschitz constant). Typical empirical values are within $v \in \mathbb{R}^{d_y \times d}$6. The output penalty $v \in \mathbb{R}^{d_y \times d}$7 is most sensitive; excessive values can slow or destabilize convergence.
• Constraint Relaxation Parameter $v \in \mathbb{R}^{d_y \times d}$8: Governs the strictness of constraint satisfaction at early timesteps. Moderate values ($v \in \mathbb{R}^{d_y \times d}$9) balance convergence speed and solution quality.
• Proximal Parameter $c \in \mathbb{R}^{d_y}$0: Must ensure $c \in \mathbb{R}^{d_y}$1 for quadratic subproblem convexity; choose $c \in \mathbb{R}^{d_y}$2.
• Activation Function Approximation: For $c \in \mathbb{R}^{d_y}$3 activations and small pre-activations, $c \in \mathbb{R}^{d_y}$4 yields closed-form updates for the $c \in \mathbb{R}^{d_y}$5- and $c \in \mathbb{R}^{d_y}$6-blocks.
• Computational Complexity: Each step involves $c \in \mathbb{R}^{d_y}$7 small matrix problems with no backpropagation through time; memory and computation cost scale linearly with sequence length.
• Initialization: Any standard random initialization is admissible. Global convergence results hold irrespective of the scale of the initial recurrent weights.

ADMMiRNN implements stable, gradient-free, block-wise ADMM training for RNNs, circumventing gradient pathologies, offering provable convergence, and supporting efficient parallelization (Tang et al., 2020).