Subspace Optimization

Latest papers

12 papers2 papers in the last four weeks, against 1 the four weeks before. 0.0% of all new papers.

Oct 14, 2026

Memory-Efficient Subspace Optimization with Residual Correction: From Momentum SGD to Muon

Memory-efficient subspace optimization has emerged as a practical approach to reducing the memory cost of LLM training while preserving optimization performance. Previous algorithms typically rely on the leading singular subspace of stochastic gradients, yet the gradient noise may dominate the subspace in later training stages and can prevent convergence. To overcome this limitation, we adopt a residual correction framework, which maintains the optimizer states in the subspace and adds the orthogonal residual gradient to the parameter update. We first integrate this scheme into momentum SGD to develop SMRC, which attains an

O(T^{-1/4})

convergence rate for a fixed orthonormal basis, matching that of full-space momentum SGD. Our analysis introduces a subspace virtual iterate to resolve the coupling between the momentum and the residual correction. Furthermore, we extend SMRC to a multi-stage variant which periodically recomputes bases to track the evolving dominant subspace and retains the same rate. Since each basis refresh breaks the consistency of the virtual iterates, we devise a boundary correction approach that aligns them across stages. Beyond momentum SGD, we incorporate our framework into Muon and propose PMRC, which applies the matrix sign to the projected momentum. Through a weighted momentum tracking analysis that accounts for the residual gradient, we establish an

O(T^{-1/4})

convergence rate for PMRC, which matches the rate of full-space Muon. We also design a multi-stage PMRC with periodically refreshed bases, which maintains a convergence rate of the same order. Extensive experimental results confirm the effectiveness and efficiency of our algorithms.

Sifan Yang, Hongyi Tao, Dingzhi Yu +5

State Key Laboratory of Novel Software Technology, Nanjing University · School of Artificial Intelligence, Nanjing University · Alibaba Token Hub, Alibaba Group +1

Subspace Optimization Memory-Efficient Training Stochastic Gradient Descent

Ask PDF

Sep 28, 2026

AIM-ZO: Activation-Informed Subspace Maintenance for Zeroth-Order LLM Fine-Tuning

Zeroth-order (ZO) optimization offers a memory-efficient alternative for LLM fine-tuning by estimating updates only from forward evaluations of perturbed parameters, without backpropagation or activation storage. However, in billion-parameter LLMs, isotropic perturbations often waste many forward evaluations on weakly informative directions. To make these evaluations more informative, existing ZO methods restrict perturbations to low-dimensional subspaces. Yet the quality of these subspaces is critical: overly compressed or poorly maintained spaces can miss useful update directions. To obtain a high-quality subspace for ZO updates, this paper proposes AIM-ZO, a ZO fine-tuning method based on Activation-Informed Subspace Maintenance. AIM-ZO uses forward activations as local directional information and continuously integrates them into a broad, evolving subspace over training. To access broader gradient-relevant structure while keeping individual perturbations low-dimensional, AIM-ZO activates only a smaller set of shared and sampled directions, decoupling the maintained width from the active width. We evaluate AIM-ZO across 5 LLMs and 11 downstream tasks under matched forward-evaluation budgets; its six-task average exceeds the strongest fully evaluated ZO baseline by 1.26 percentage points on OPT-2.7B and MeZO by 2.85 percentage points on OPT-30B. Our code is available at https://github.com/EkkoXy/AIM-ZO

Yue Xie, Zhi Zheng, Yunpeng Ba +5

Southern University of Science and Technology · National University of Singapore · Huawei Technologies Ltd. +1

Subspace Optimization Memory-Efficient Fine-Tuning Zeroth-Order Optimization LLM Fine-Tuning

Ask PDF

Sep 16, 2026

Gradient Descent with Stochastic Subspaces via Persistence of Memory

Stochastic subspace methods have gained popularity as gradient descent based techniques for large scale optimisation problems, especially in distributed settings. In this paper, we introduce the technique of "persistence of memory" to greatly extend and improve the random subspace methods. To this end, we leverage a vector that is only weakly correlated with the gradient in order to provide a guiding structure to the generative process of the random subspace along which the descent is going to take place. This guidance vector may be fixed for a large number of iterations, only to be refreshed at wide intervals (on whose size we can provide guarantees in terms of problem parameters). In important machine learning settings, such as optimisation problems embodying sparsity or a minibatch structure, we show that the guidance vector can be obtained in an effective and computationally inexpensive manner by leveraging the structured properties of the problem. En route, we establish to our knowledge the first theoretical analysis of classical SSD methods for sparse functions. In a local neighbourhood of the optimum, we demonstrate an alignment phenomenon of our gradient estimates with a low-lying eigenvector of the Hessian, allowing a once-for-all computation of the guidance vector which renders the method computationally favourable even in scenarios with unstructured objectives.

Subhroshekhar Ghosh, Clement Z. Q. Ng, Pierre-Louis Poirion +1

Department of Mathematics, National University of Singapore, Singapore · Center for Advanced Intelligence Project, RIKEN, Tokyo, Japan · Department of Mathematical Informatics, The University of Tokyo, Tokyo, Japan

Gradient Descent Stochastic Optimization Convergence Subspace Optimization

Ask PDF

Sep 5, 2026

FedSubMuon: Communication-Efficient Federated LLM Fine-Tuning via Structured Subspace Muon

Federated fine-tuning adapts large language models (LLMs) to decentralized client data, but its scalability in cross-device training is often limited by the high communication cost. Muon is an optimizer that improves optimization performance by orthogonalizing momentum for matrix-valued parameters. Existing federated Muon methods demonstrate the benefit of matrix-aware optimization in federated learning, but still require transmitting full layer-size updates and optimizer state. A natural way to reduce communication is to directly apply Muon to LoRA factors, but this changes the optimized object and weakens Muon's matrix-aware update geometry. We propose FedSubMuon, a communication-efficient federated Muon fine-tuning method that optimizes compact coefficient matrices within shared structured subspaces. This design keeps Muon on a single matrix-valued trainable object, while reducing the client upload to compact coefficient matrices. We further introduce FedSubMuon-GT, an accuracy-oriented extension that uses projected gradients to adapt tracked subspace bases toward task-relevant gradient directions. Experiments on instruction tuning and mathematical reasoning show that FedSubMuon-GT achieves the best overall accuracy on four of five dataset-model pairs, while FedSubMuon performs best under all matched communication budgets. On Dolly-15K, the closest communication baseline requires 5.5 times and 1.4 times more total communication on Llama-1B and Qwen-4B, respectively.

Shaolong Chen, Youming Tao, Shuzhen Chen +3

King Abdullah University of Science and Technology (KAUST) · Provable Responsible AI and Data Analytics (PRADA) Lab · Imperial College London +3

Subspace Optimization Federated Fine-Tuning LLM Fine-Tuning Federated Learning Muon Optimizer

Ask PDF

Aug 12, 2026

Predicting When Random Low-Dimensional Reparameterizations Train Neural Networks

Neural networks can often be trained or fine-tuned through random low-dimensional reparameterization, where a small latent vector is mapped into a full parameter update by a frozen random map. This raises a practical question: how large must the latent search space be to reach a low-loss region? We first express the known accessibility transition in an equivalent conic form, centered for compact convex targets at the statistical dimension of the polar cone. Our main theoretical contribution is an orientation-resolved quadratic master formula that predicts the random-slice residual from both the curvature spectrum and the reference-to-solution displacement profile. It yields a self-consistent isotropic-orientation predictor and, in a conservative radius-only specialization, recovers the earlier Gaussian-width quadratic bound. Building on this analysis, we introduce Random Mapping Networks (RaMaN), which instantiate the predicted latent dimension using structured Hadamard or seed-regenerated Gaussian maps. These constructions avoid the O(dP) storage of dense random maps and reduce optimizer-state memory from O(P) to O(d). We also develop matrix-free curvature approximations and sweep-free dimension selection. Across controlled quadratic and neural-curvature experiments, the orientation-resolved predictor closely tracks measured transition locations and outperforms orientation-agnostic approximations when displacement direction matters. End-to-end experiments further show sharp, protocol-dependent training transitions across image and language models.

Andrew Cheng, Ali Eslamian, Jie Cheng +2

Department of Computer Science, Tsinghua University, Beijing, China · Department of Computer Science, The University of Manchester, Manchester, UK · Department of Computer Science, University of Kentucky, Lexington, KY, USA +3

Subspace Optimization

Ask PDF

Jul 15, 2026

Dysco: Dynamic Subspace Boosting to Mitigate LoRA Interference in Federated Learning

Federated fine-tuning of large pre-trained models increasingly relies on Low-Rank Adaptation (LoRA) to reduce communication and computation, but heterogeneous clients can make adapter aggregation unstable. We identify the data-parameter interference as a geometric source of this instability. This interference is controlled by the alignment between LoRA update subspaces and client activations, suggesting that federated LoRA aggregation should be viewed not only as parameter averaging but also as subspace allocation. We propose Dynamic Subspace Boosting (Dysco), a plug-in method that allocates client-specific LoRA subspaces in a federated and dynamic manner. In each round, clients compute activation-insensitive subspaces from local representations and transmit only the resulting bases; the server then constructs client-specific merged subspaces through a closed-form solution that maximizes compatibility with other clients' insensitive directions. To handle representation drift, Dysco performs multi-round subspace boosting to preserve past update directions while adapting to future representations. We provide a convergence analysis that embeds the data-parameter interference as an aggregation-error term in a standard federated optimization bound, and prove that Dysco's server-fixed merged subspaces yield a tighter upper bound on this error. Experiments on controlled synthetic federated tasks and on MIMIC-IV clinical-note classification with Llama-3.2-1B show that Dysco substantially reduces interference, reduces the final-round synthetic training loss by up to 9 times relative to baselines under the orthogonal-subspace partition the theory identifies, improves all five tested FL algorithms by up to 4.3% on MIMIC, outperforms recent federated LoRA methods, and adds only 0.9% wall-clock overhead. Our code is available at https://github.com/illidanlab/Dysco.

Haobo Zhang, Jiankun Wang, Suraj Rajendran +5

University of Michigan · Cornell University · University of Pennsylvania

Subspace Optimization Federated Learning Aggregation Low-Rank Adaptation Federated Fine-Tuning Federated Learning Heterogeneous Federated Learning

Ask PDF

May 13, 2026

Stochastic global optimization of continuous functions via random walks on Grassmannians

We introduce a stochastic global optimization method based on random walks on Grassmannian manifolds. To minimize a continuous objective

\ell:\mathbb{R}^d\rightarrow\mathbb{R}

, the method repeatedly samples random

k

-dimensional linear subspaces (with

k\ll d

), solves the resulting low-dimensional restrictions of these problems to these subspaces using an arbitrary black-box optimizer, and updates the iterate (which monotonically improves upon the previous iterate). Unlike classical optimization analyses that rely on convexity, smoothness, Lipschitz bounds, or Polyak-Lojasiewicz-type conditions, our convergence guarantees depend only on the geometric distribution of restricted minima across the

k

-dimensional subspaces passing through a given point in

\mathbb{R}^d

. We identify a gap parameter -- an analogue of a spectral gap for random walks -- that controls the rate at which the iterates approach the global minimum value. Finally, we argue that the same analysis yields a blind-spot robustness property: sufficiently narrow, deep dips of the loss function (small-measure regions where

\ell

spikes downward) have limited influence on the algorithm's trajectory, since they are unlikely to be encountered by random subspace sampling.

Kartik Gupta, Stephen D. Miller, Pradeep Ravikumar +1

Stochastic Optimization Subspace Optimization Black-Box Optimization Global Optimization Stochastic Optimization Convergence

Ask PDF

May 11, 2026

BROS: Bias-Corrected Randomized Subspaces for Memory-Efficient Single-Loop Bilevel Optimization

Stochastic bilevel optimization (SBO) has become a standard framework for hyperparameter learning, data reweighting, representation learning, and data-mixture optimization in deep learning. Existing exact single-loop SBO methods and memory-efficient surrogate SBO methods either create severe memory pressure for large lower-level neural networks or lack competitive convergence guarantees under standard assumptions. In this paper, we propose BROS, a memory-efficient single-loop SBO method with the same convergence rate order as exact single-loop SBO methods. BROS performs lower and auxiliary updates in randomized subspaces with a Rademacher bi-probe correction that recovers an unbiased Hessian-action estimator. We prove that BROS preserves the

\mathcal O(\varepsilon^{-2})

sample complexity of MA-SOBA for finding an

\varepsilon

-stationary point under only standard assumptions. Experiments on hyper-data cleaning, data-mixture learning, hyper-representation learning, and ViT sample reweighting show that BROS reduces peak memory by up to 44.9% while closely matching full-space baseline performance.

Hengrui Zhang, Boao Kong, Engao Zhang +1

1Sichuan University · 2Peking University

Stochastic Bilevel Optimization Bilevel Optimization Hyperparameter Optimization Subspace Optimization

Ask PDF

May 1, 2026

Randomized Subspace Nesterov Accelerated Gradient

Randomized-subspace methods reduce the cost of first-order optimization by using only low-dimensional projected-gradient information, a feature that is attractive in forward-mode automatic differentiation and communication-limited settings. While Nesterov acceleration is well understood for full-gradient and coordinate-based methods, obtaining accelerated methods for general subspace sketches that use only projected-gradient information and can improve over full-dimensional Nesterov acceleration in oracle complexity is technically nontrivial. We develop randomized-subspace Nesterov accelerated gradient methods for smooth convex and smooth strongly convex optimization under matrix smoothness and generic sketch moment assumptions. The key technical ingredient is a three-sequence formulation tailored to matrix smoothness, which recovers the corresponding classical Nesterov methods in the full-dimensional case. The resulting theory establishes accelerated oracle-complexity guarantees and makes explicit how matrix smoothness and the sketch distribution enter the complexity. It also provides a unified basis for comparing sketch families and identifying when randomized-subspace acceleration improves over full-dimensional Nesterov acceleration in oracle complexity.

Gaku Omiya, Pierre-Louis Poirion, Akiko Takeda

Department of Mathematical Informatics, The University of Tokyo, Tokyo, Japan · Center for Advanced Intelligence Project, RIKEN, Tokyo, Japan

Gradient Descent Randomized Sketching Accelerated Gradient Methods Subspace Optimization

Ask PDF

Apr 28, 2026

Subspace Optimization for Efficient Federated Learning under Heterogeneous Data

Federated learning increasingly operates in a large-model regime where communication, memory, and computation are all scarce. Typically, non-IID client data induce drift that degrades the stability and performance of local training. Existing remedies such as SCAFFOLD introduce heterogeneity-correction mechanisms to address this challenge, but they incur substantial extra communication and memory overhead. This paper proposes a subspace optimization method for federated learning (SSF), which performs heterogeneity-corrected optimization in a low-dimensional subspace using only projected quantities, while preserving full-dimensional control information through a backfill-style update that retains residual components whenever the active subspace changes. Under standard smoothness and bounded-variance assumptions, SSF attains a non-asymptotic rate of order

\widetilde{\mathcal{O}}(1/T+1/\sqrt{NKT})

. Experiments show favorable accuracy--efficiency trade-offs under heterogeneous data.

Shuchen Zhu, Zhengyang Huang, Yuqi Xu +1

Peking University · Beihang University

Subspace Optimization Non-IID Federated Learning Federated Learning Heterogeneous Federated Learning

Ask PDF

Mar 30, 2026

Subspace Optimization for Backpropagation-Free Continual Test-Time Adaptation

We introduce PACE, a backpropagation-free continual test-time adaptation system that directly optimizes the affine parameters of normalization layers. Existing derivative-free approaches struggle to balance runtime efficiency with learning capacity, as they either restrict updates to input prompts or require continuous, resource-intensive adaptation regardless of domain stability. To address these limitations, PACE leverages the Covariance Matrix Adaptation Evolution Strategy with the Fastfood projection to optimize high-dimensional affine parameters within a low-dimensional subspace, leading to superior adaptive performance. Furthermore, we enhance the runtime efficiency by incorporating an adaptation stopping criterion and a domain-specialized vector bank to eliminate redundant computation. Our framework achieves state-of-the-art accuracy across multiple benchmarks under continual distribution shifts, reducing runtime by over 50% compared to existing backpropagation-free methods.

Damian Sójka, Sebastian Cygert, Marc Masana

Poznan University of Technology · Akces NCBR · NASK - National Research Institute +3

Subspace Optimization Continual Test-Time Adaptation Test-Time Adaptation Derivative-Free Optimization

Ask PDF

Jan 20, 2026

Linearized subspace refinement framework to expose hidden accuracy in trained neural networks

Neural networks trained by gradient-based methods often exhibit optimization-induced accuracy plateaus in scientific machine learning tasks. We present Linearized Subspace Refinement (LSR), an architecture-agnostic post-training framework that exploits the local linearized model at a fixed trained state. By solving a reduced direct least-squares problem in a Jacobian-defined low-dimensional space, LSR computes a subspace-optimal linearized correction and yields a refined predictor with markedly improved accuracy. Across function approximation, data-driven operator learning, physics-informed operator fine-tuning, and noisy inverse problems, LSR shows that standard nonlinear training can remain far above this subspace-attainable error level. Similar accuracy plateaus persist even for the convex quadratic problem from local linearization when solved with standard iterative optimizers, identifying numerical ill-conditioning as a primary bottleneck. LSR frequently delivers order-of-magnitude error reductions, while the subspace rank provides an explicit capacity-control mechanism that balances correction strength, numerical stability, and noise sensitivity. Together, LSR exposes conditioning-limited attainable accuracy in trained-state linearized models and provides direct access to it.

Wenbo Cao, Weiwei Zhang

Institute of AI for Industries, Chinese Academy of Sciences, Nanjing 211135, China · Institute of Computing Technology, Chinese Academy of Sciences, Beijing 100190, China · School of Aeronautics, Northwestern Polytechnical University, Xi’an 710072, China +2

Subspace Optimization Neural Network Optimization Neural Network Approximation Theory Neural Network Training Dynamics

Ask PDF