Table of Contents
Fetching ...

SALAAD: Sparse And Low-Rank Adaptation via ADMM

Hao Ma, Melis Ilayda Bal, Liang Zhang, Bingcong Li, Niao He, Melanie Zeilinger, Michael Muehlebach

TL;DR

SALAAD tackles the challenge of deploying large language models under heterogeneous compute and memory budgets by inducing sparse and low‑rank structure during pretraining. It combines an ADMM‑based decomposition with an adaptive I‑controller and a homomorphic parameter allocation (HPA) strategy to produce a structured surrogate that supports continuous, deployment‑aware capacity scaling without architectural modification. The approach achieves competitive perplexity across model scales while enabling elastic deployment across memory budgets and enabling post hoc compression via a global budget with minimal tuning. This work demonstrates that embedding layers can exhibit stable SLR under adaptive induction, and provides a principled, plug‑and‑play path toward memory‑efficient LLMs with predictable performance‑capacity trade‑offs.

Abstract

Modern large language models are increasingly deployed under compute and memory constraints, making flexible control of model capacity a central challenge. While sparse and low-rank structures naturally trade off capacity and performance, existing approaches often rely on heuristic designs that ignore layer and matrix heterogeneity or require model-specific architectural modifications. We propose SALAAD, a plug-and-play framework applicable to different model architectures that induces sparse and low-rank structures during training. By formulating structured weight learning under an augmented Lagrangian framework and introducing an adaptive controller that dynamically balances the training loss and structural constraints, SALAAD preserves the stability of standard training dynamics while enabling explicit control over the evolution of effective model capacity during training. Experiments across model scales show that SALAAD substantially reduces memory consumption during deployment while achieving performance comparable to ad-hoc methods. Moreover, a single training run yields a continuous spectrum of model capacities, enabling smooth and elastic deployment across diverse memory budgets without the need for retraining.

SALAAD: Sparse And Low-Rank Adaptation via ADMM

TL;DR

SALAAD tackles the challenge of deploying large language models under heterogeneous compute and memory budgets by inducing sparse and low‑rank structure during pretraining. It combines an ADMM‑based decomposition with an adaptive I‑controller and a homomorphic parameter allocation (HPA) strategy to produce a structured surrogate that supports continuous, deployment‑aware capacity scaling without architectural modification. The approach achieves competitive perplexity across model scales while enabling elastic deployment across memory budgets and enabling post hoc compression via a global budget with minimal tuning. This work demonstrates that embedding layers can exhibit stable SLR under adaptive induction, and provides a principled, plug‑and‑play path toward memory‑efficient LLMs with predictable performance‑capacity trade‑offs.

Abstract

Modern large language models are increasingly deployed under compute and memory constraints, making flexible control of model capacity a central challenge. While sparse and low-rank structures naturally trade off capacity and performance, existing approaches often rely on heuristic designs that ignore layer and matrix heterogeneity or require model-specific architectural modifications. We propose SALAAD, a plug-and-play framework applicable to different model architectures that induces sparse and low-rank structures during training. By formulating structured weight learning under an augmented Lagrangian framework and introducing an adaptive controller that dynamically balances the training loss and structural constraints, SALAAD preserves the stability of standard training dynamics while enabling explicit control over the evolution of effective model capacity during training. Experiments across model scales show that SALAAD substantially reduces memory consumption during deployment while achieving performance comparable to ad-hoc methods. Moreover, a single training run yields a continuous spectrum of model capacities, enabling smooth and elastic deployment across diverse memory budgets without the need for retraining.
Paper Structure (18 sections, 16 equations, 8 figures, 8 tables, 1 algorithm)

This paper contains 18 sections, 16 equations, 8 figures, 8 tables, 1 algorithm.

Figures (8)

  • Figure 1: Comparison of SALAAD training with and without embedding layer inclusion on a LLaMA-based 350M model. (a) Training loss trajectories. (b) Convergence of effective rank ratio and density in the embedding layer. (c) Convergence behavior of a randomly selected Transformer block. (d) Singular value spectra of the learned low-rank components. Overall, the results indicate that including the embedding layer does not affect training dynamics.
  • Figure 2: Perplexity versus parameter count for SALAAD models under different parameter budgets, compared with vanilla models. All results are obtained by applying HPA strategy.
  • Figure 3: Effect of allocation ratio $\kappa$ on model performance under different parameter budgets for LLaMA-based (a) 130M model, (b) 350M model, and (c) 1B model. The gray region indicates the relatively stable range of optimal allocation ratio $\kappa^{\star}$ across different budgets.
  • Figure 4: Post-hoc RPCA results on standard-trained models. (a) Results for the LLaMA 1B model. (b) Results for the LLaMA 3.2 3B model. For each model, representative layers from shallow, middle, and deep regions are selected, and for each layer the effective rank ratios and sparsity levels obtained after RPCA decomposition are reported for different matrix types.
  • Figure 5: Post-hoc RPCA results on SALAAD-trained LLaMA-based 1B model. For several representative layers from shallow, middle, and deep regions, the effective rank ratios and sparsity levels obtained after RPCA decomposition are compared with the ground-truth values from the original SLR components learned by SALAAD. Black and gray boxes denote the recovered effective rank ratios and sparsity levels, respectively, while red and light-red boxes denote the corresponding ground-truth values.
  • ...and 3 more figures

Theorems & Definitions (1)

  • Definition 4.1: Effective Rank Ratio under Energy Coverage