Preconditioned Inexact Stochastic ADMM for Deep Models (PISA / SISA / NSISA)

TL;DR

PISA reframes neural-network training as a stochastic ADMM problem: split the dataset into m batches, attach an auxiliary variable and a Lagrange multiplier to each, and solve each subproblem inexactly with one preconditioned gradient step. The framework supports arbitrary preconditioning matrices, yielding two named variants — SISA (second-moment preconditioning, like Adam) and NSISA (Newton–Schulz orthogonalized momentum, like Muon). The headline claims are theoretical (convergence under only Lipschitz-continuity of the gradient — no bounded gradients, no bounded variance, no IID assumption, no unbiased estimator) and empirical (matches or beats AdamW / Muon / Shampoo / SOAP / Adam-mini across GPT2-Nano→XL, ResNet/VGG/DenseNet on CIFAR-10/ImageNet, and WGAN/WGAN-GP). The non-IID claim is the most distinctive: under extreme label skew (one class per federated client) SISA hits ~95% on MNIST where FedAvg / FedProx / Scaffold plateau at ~54%.

Key claims

PISA converges under a single assumption — Lipschitz continuity of the gradient on a bounded region — and explicitly drops the bounded-gradient, bounded-second-moment, bounded-variance, unbiased-estimator, and IID assumptions that every other compared stochastic method requires [Table 1, §Contributions / “Strong convergence theory under a sole assumption”].
The algorithmic core is preconditioned inexact stochastic ADMM: per-batch auxiliary variable w_i + multiplier π_i, with a (ρ_i/2)⟨w_i − w^{ℓ+1}, Q_i^{ℓ+1}(w_i − w^{ℓ+1})⟩ proximal term whose matrix Q_i is the preconditioner [Eq. (6a-c), Algorithm 1, Remark 1].
SISA = PISA with diagonal second-moment preconditioner (analogue of Adam’s v_t); NSISA = PISA with Newton–Schulz-orthogonalized momentum preconditioner (analogue of Muon) [§Contributions / “A general algorithmic structure”].
Under extreme federated label skew (1 label per client) on MNIST, SISA reaches 94.97% test accuracy versus a best-of-baselines (FedAvg / FedProx / Scaffold / …) of 54.33%, while using one local update per aggregation step against the baselines’ multiple local updates [§Results / “Data heterogeneity”, Table 2].
On CIFAR-10 across VGG-11, ResNet-34, DenseNet-121, SISA achieves the highest test accuracy on ResNet-34 and DenseNet-121 and matches the best baseline (SGD-M) on VGG-11; on ImageNet/ResNet-18 SISA beats most adaptive baselines but trails SGD-M and AdaBelief [Table 3, §Results / “Classification by vision models”].
Training GPT2 (Nano 125M → Medium 330M → XL 1.5B) on FineWeb, NSISA’s validation-loss gap over AdamW / Muon / Shampoo / SOAP / Adam-mini widens with scale, and NSISA’s parallel implementation gives a clear wall-clock advantage at GPT2-XL [Fig. 2, §Results / “Training LLMs”].
On WGAN and WGAN-GP / CIFAR-10, SISA achieves both the fastest FID convergence and the lowest final training and test FID across 5 independent runs versus Adam, RMSProp, AdaBelief — directly addressing the optimizer-stability problem in GAN training [Fig. 3, §Results / “Generation tasks by GANs”].
Algorithm structure is inherently parallel: subproblems on the m batches are independent given the current global w, mapping cleanly onto centralized federated learning where the server aggregates client-side (w_i, π_i) updates [Algorithm 2, §Methods].

Method

The training objective (1/m)Σ_i (1/|D_i|) Σ_{x∈D_i} f(w; x) + (μ/2)‖w‖² is rewritten in equivalent split form by introducing per-batch auxiliaries w_i constrained to equal the global w, with multipliers π_i and penalty constants σ_i. The augmented Lagrangian is then minimized by alternating three updates per iteration ℓ: (a) closed-form update of global w^{ℓ+1} averaging w_i^ℓ − π_i^ℓ/σ_i^ℓ; (b) an inexact update of each w_i^{ℓ+1} via a single preconditioned stochastic-gradient step taking the stochastic gradient ∇f_i(w_i^ℓ; ξ_i^ℓ) and the preconditioner Q_i^{ℓ+1}; (c) multiplier update π_i^{ℓ+1} = π_i^ℓ + σ_i^{ℓ+1}(w_i^{ℓ+1} − w^{ℓ+1}). The proximal term (ρ_i/2)⟨w_i − w^{ℓ+1}, Q_i^{ℓ+1}(w_i − w^{ℓ+1})⟩ is the slot where SGD-family preconditioners are dropped in — diagonal second-moment (SISA, ≈ AdamW-style), Hessian-based (full second order), or Newton–Schulz orthogonalized momentum (NSISA, ≈ Muon-style). Convergence is proved under Lipschitz continuity of the gradient on a bounded region (condition 12 in the paper), which is weaker than global Lipschitz continuity and satisfied by any twice-continuously-differentiable loss; bounded-gradient, bounded-variance, IID-sampling, and unbiased-estimator assumptions are all dropped from the proof. Source code is at https://github.com/Tracy-Wang7/PISA. GPT2 experiments run on 4×H100-80GB; everything else on a single H100-80GB.

Results

Federated MNIST, 1-label skew (per-client homogeneous class): SISA 94.97% vs best baseline 54.33% [Table 2]. Also reported strong on CIFAR-10, FMNIST, Adult under 1/2/3-label skew.
CIFAR-10 single-host: SISA best on ResNet-34 and DenseNet-121; matches SGD-M on VGG-11 [Table 3].
ImageNet / ResNet-18: SISA outperforms most adaptive baselines, trails SGD-M and AdaBelief [Table 3].
GPT2-Nano / Medium / XL on FineWeb: NSISA validation-loss gap over AdamW, Muon, Shampoo, SOAP, Adam-mini widens as N grows from 125M → 330M → 1.5B; NSISA also wins on wall-clock at XL via parallel implementation [Fig. 2]. Memory overhead numbers are in the supplement.
WGAN and WGAN-GP / CIFAR-10: SISA lowest training FID and lowest test FID (computed on 64,000 generated images) across 5 runs vs Adam, RMSProp, AdaBelief; also fastest FID convergence [Fig. 3].

Why it’s interesting

The PISA framework lands directly in the optimizer-side fork of the training-stability conversation traced in Training stability at scale — alongside Muon, Shampoo, and the spectral-sphere optimizer of Controlled LLM Training on Spectral Sphere. Both PISA and SSO take the position “stability lives in the optimizer,” but the recipe is different: SSO retracts onto a per-module spectral sphere (constrained geometry); PISA decomposes the loss into per-batch subproblems and absorbs known preconditioners (Adam-style second moment, Muon-style Newton–Schulz orthogonalization) into an ADMM proximal slot. The most interesting empirical claim for a Luma-style team is the increasing validation-loss gap over Muon/Shampoo/SOAP as GPT2 scales from 125M to 1.5B — if that trend holds at >10B, NSISA becomes a serious AdamW/Muon replacement candidate. The non-IID federated result is also worth flagging for any team that runs multi-node training where batches are not shuffled across the cluster (a regime the paper argues breaks every other adaptive optimizer’s convergence theorem in principle, even if it works in practice). One caveat: GAN training stability has been “solved” by many published optimizers in turn; FID gains on WGAN/CIFAR-10 are a low bar relative to the modern diffusion-vs-GAN baselines that aren’t in this evaluation.