FlexRank: Nested Low-Rank Knowledge Decomposition for Adaptive Model Deployment

Conference: ICML 2026
arXiv: 2602.02680
Code: https://github.com/RickZack/FlexRank
Area: Model Compression / Elastic Inference / Low-Rank Decomposition
Keywords: Elastic models, low-rank decomposition, nested submodels, knowledge distillation, Pareto frontier

TL;DR

FlexRank performs activation-aware low-rank decomposition (DataSVD) on each linear layer of pre-trained large models and uses dynamic programming to select a set of strictly nested submodels corresponding to different compute budgets in \(O(L\cdot K)\) time. These submodels are jointly trained via knowledge distillation with shared weights. Finally, Gauge-Aligned Reparametrization is used to translate rank savings into actual FLOPs savings—obtaining a "family" of deployable models for LLMs and ViTs that approach the true Pareto frontier with a single training run.

Background & Motivation

Background: LLMs and ViTs have expanded to billions of parameters, making training from scratch affordable only for a few institutions. The mainstream approach in the community is to reuse pre-trained weights for downstream adaptation with PEFT (e.g., LoRA) or to reduce deployment costs using quantization/pruning.

Limitations of Prior Work: PEFT only modifies a small portion of parameters, leaving the compute structure of the backbone unchanged, thus deployment costs remain "one-size-fits-all." While quantization and pruning reduce computation, quantization-aware training requires pipeline modifications, and structural sparsity depends on hardware kernels. More importantly, these methods typically produce a single compression ratio; different hardware (e.g., a phone vs. a server) requires repeated retraining or maintaining multiple sets of weights.

Key Challenge: Existing elastic solutions follow two paths: (i) Post-hoc subnetwork selection (PTS) after training a full model—Theorem 4.1 proves the probability of obtaining a Pareto optimal submodel this way is "zero"; or (ii) Joint training of all submodels (ASL), where all subnets compete for the same representation capacity. Theorem 4.2 proves the suboptimality gap for each rank in ASL is strictly greater than zero. Neither path yields a family of submodels truly close to the Pareto frontier.

Goal: Starting from a pre-trained model, construct a set of shared weights \(\theta\) and a sequence of strictly nested masks \(\mathbf{m}_1 \preceq \mathbf{m}_2 \preceq \dots \preceq \mathbf{m}_K\), such that submodels extracted under \(K\) different compute budgets \(\beta_k\) simultaneously approach the true Pareto frontier.

Key Insight: The authors observe that SVD provides a natural "importance order" for weights in each layer (singular values from largest to smallest). The seemingly more restrictive constraint of nesting actually prevents the mutual interference seen in ASL: the \((r+1)\)-th column only needs to learn the residual \(A_{r+1}-A_r\) between the \((r+1)\)-th and \(r\)-th rank SVD truncations, avoiding competition for capacity with smaller submodels.

Core Idea: Use layer-wise activation-aware SVD to determine local importance, employ DP to aggregate local orders into a global nested submodel family, use distillation to transform "independent layer decompositions" into a "collaborative end-to-end elastic model," and finally use gauge reparametrization during inference to convert \(r\)-th rank truncation into \(\mathcal{O}((m+n-r)r)\) FLOPs.

Method

Overall Architecture

The input consists of a pre-trained model \(f(\cdot; \theta_{\mathrm{orig}})\), a small calibration set \(\mathcal{Z}\) (approx. \(10^3\) samples), and a set of target budgets \(\mathcal{B}=\{\beta_k\}_{k=1}^K\). The output is a single set of shared parameters \(\theta=\{(U_l, V_l)\}_{l=1}^L\) and a sequence of nested masks \(\mathcal{M}^\star=\{\mathbf{m}_k^\star\}\). At deployment, given a budget \(\beta\), the corresponding submodel \(\theta_\beta\) can be assembled in \(O(L)\) time without additional training.

The pipeline consists of three stages: (1) Layer Decomposition — Activation-aware SVD is performed independently for each layer to obtain ordered factors \((U_l, V_l)\); (2) Nested Submodel Search — Under the additive-error assumption, DP compresses the \(K^L\) possible combinations of \(K\) ranks per layer into \(O(L\cdot K)\); (3) Knowledge Distillation Consolidation — Shared weights are trained by sampling from the \(K\) nested masks and distilling from teacher logits.

Key Designs

1. DataSVD: Activation-aware Layer Decomposition:

   • Function: Decomposes each layer weight \(W_l\) into \(U_l V_l^\top\). The objective is not to minimize \(\|W_l - U_l V_l^\top\|_F^2\), but to minimize the output error \(\mathbb{E}_{\mathbf{x}_l}\bigl[\|(W_l-U_l V_l^\top)\mathbf{x}_l\|_2^2\bigr]\). Thus, singular directions are determined by the activation covariance "truly important to the task" rather than the weights themselves.
   • Mechanism: Captures activation matrices \(\mathbf{X}_l\) using the calibration set and performs closed-form SVD on the weighted problem. The authors prove implementation can achieve space complexity \(\mathcal{O}(n_l^2)\) independent of sample size \(N\). This step serves as initialization, providing the "importance order" for DP (Remark 3.1 notes SVD alone is insufficient and requires distillation).
   • Design Motivation: Pure SVD often leads to severe performance drops in LLMs (Fig. 4 shows collapse at 20% compression) because large weights do not necessarily represent large contributions to real inputs. Activation-awareness aligns "important directions" with the "real input distribution."

2. Nested Submodel Search + Dynamic Programming:

   • Function: Selects \(K\) strictly nested submodels \(\mathbf{m}_1 \preceq \dots \preceq \mathbf{m}_K\) from \(K^L\) possible rank combinations to approximate the Pareto optimum for each budget \(\beta_k\).
   • Mechanism: Enumerates \(K\) candidate ranks for each layer \(l\), calculating the cost reduction \(\Delta c\) and error increase \(\Delta e\) from truncation to obtain a local Pareto table \(\mathcal{Q}_l\). Assuming additivity of layer errors, DPRankSelection combines these local tables into a global nested mask sequence in \(\mathcal{O}(L\cdot K)\) time. The additivity assumption's ranking fidelity was validated in small-scale enumerable settings.
   • Design Motivation: The nesting constraint is derived theoretically: Thm 4.1 says the probability of PTS finding Pareto optimality is 0; Thm 4.2 says ASL leaves a suboptimality gap of at least \(\frac{1}{k}(r\lambda-\sum_{i\le r}\sigma_i)^2\); Thm 4.3 proves nested training (NSL) makes the gap exactly 0 for each rank because the \((r+1)\)-th column only learns the residual. This is the core theoretical contribution.

3. Gauge-Aligned Reparametrization (GAR):

   • Function: Further reduces actual inference cost from \(\mathcal{O}(mr+nr)\) to \(\mathcal{O}((m+n-r)r)\) for rank-\(r\) truncation, making it strictly cheaper than dense \(\mathcal{O}(mn)\) and ensuring \(r\) reduction translates linearly into FLOPs reduction.
   • Mechanism: Exploits the non-uniqueness of \(UV^\top\) decomposition. By introducing a gauge \(G=U_{1:r,:}^{-1}\), one gets \(UV^\top = (UG)(G^{-1}V^\top) = \tilde{U}\tilde{V}^\top\), where the top \(r\times r\) block of \(\tilde{U}\) is aligned to \(I_r\). Thus, the top part requires neither storage nor multiplication, leaving only \((m-r)\times r\) of \(\hat{U}\) for calculation. GAR preprocessing is a one-time \(\mathcal{O}(r^3)\) inversion.
   • Design Motivation: In the original \((U, V)\) form, FLOPs only become competitive with dense kernels when \(r \ll \min(m, n)\). GAR eliminates this threshold, ensuring immediate benefits for any \(r < \min(m, n)\).

Loss & Training

With fixed \(\mathcal{M}^\star\), a mask \(\mathbf{m}_t^\star\) is sampled with weight \(\alpha_k\) at each step. The submodel output is aligned with the original teacher \(f(\cdot; \theta_{\mathrm{orig}})\):

\[\ell_k(\theta)=\mathbb{E}_{\mathbf{d}}\bigl[\mathcal{L}_{\text{KD}}(f(\mathbf{d};\mathcal{T}_{\mathbf{m}_k^\star}(\theta)), f(\mathbf{d};\theta_{\mathrm{orig}}))\bigr]\]

The objective is \(\min_\theta \sum_k \alpha_k \ell_k(\theta)\), solved via standard gradient descent. On Llama-3.2-1B, using only 5B tokens (167× less than LayerSkip's 839B), it matches or exceeds heavier baselines.

Key Experimental Results

Main Results

Configuration	Evaluation	FlexRank	SOTA Comparison	Description
Llama-3.2-1B/3B/8B, 5B tokens	avg acc on commonsense lm-eval-harness	Consistently leads at 20–80% budgets	SVD/DataSVD collapse at 20% parameters; ACIP is outperformed at low budgets	Fig. 4-top
DINOv3 ViT-L/16 → ViT-7B/16, ImageNet-1K	top-1 acc	Remains close to full model at 30%	Baselines significantly lag at all budgets	Fig. 4-bottom
Llama-3.2-1B, LoRA fine-tuned on Math/Code	math/code avg acc	Smooth decline from base→1×→0.8×→0.4× (math: 25.7→25.0→20.5→13.6)	—	Tab. 1, submodels support downstream LoRA adaptation

Ablation Study

Configuration	Key Findings	Value
PTS (Post-hoc selection)	Pareto gap always > 0 (Thm 4.1)	Post-hoc subnet extraction is inherently suboptimal
ASL (Joint training)	Strictly positive gap \(\ge \frac{1}{k}(r\lambda-\sum\sigma_i)^2\) (Thm 4.2)	Subnets interfere and compete for capacity
NSL (Ours: Nested training)	gap = 0 (Thm 4.3)	Nesting is sufficient to recover the Pareto frontier
Independent layer training	Consistently poor performance (Fig. 7b)	End-to-end distillation is necessary for non-linear flow
DataSVD calibration samples	Saturates at 128 samples (Fig. 7a)	Minimal calibration overhead
Hierarchical compression heatmap	Middle attention layers (c_proj) pruned last (Fig. 6)	DP allocates budget based on importance variability

Key Findings

• Nesting is a "necessity" for Pareto elasticity, not a heuristic: The theoretical framework (Thm 4.1/4.2/4.3) narrows down the optimal route to "nesting + joint training."
• GAR enables "immediate FLOPs savings" for low-rank compression: By eliminating the threshold where low-rank exceeds dense compute costs, GAR is the key engineering trick to translate theoretical savings into runtime speedup.
• Amortized training cost: While training requires the full rank (approx. 2× VRAM and slower than dense forward), obtaining \(K\) deployable models from one run is much more efficient than separate retraining.

Highlights & Insights

• "Theory-first" design: The paper uses Thm 4.1/4.2 to "rule out" the two most intuitive paths (PTS and ASL) before proving Thm 4.3. This rigorous approach is more convincing than empirical "trial and error."
• Decoupled universal trick (GAR): GAR was applied to all low-rank baselines in experiments, ensuring the comparison focused on the "algorithm" rather than "engineering optimization gaps."
• Taming combinatorial explosion with DP: The additivity assumption allows for a tractable \(\mathcal{O}(LK)\) search instead of an impossible \(K^L\) search, a strategy transferable to any layer-independent scoring scenario.

Limitations & Future Work

• The additivity assumption is strictly speaking invalid for deep non-linear networks; its impact on 8B+ models beyond the validated small networks remains unquantified.
• Training requires full-rank \((U,V)\) storage (~2× VRAM), which may be prohibitive for 70B+ models without sharded factors.
• Input-adaptive routing (dynamically selecting budget per token) is not evaluated, though FlexRank naturally supports this.
• Combinations with quantization, depth-elasticity, and MoE are yet to be explored.

Related Work & Insights

• vs ACIP (Genzel et al., 2025): ACIP uses SVD+LoRA with frozen bases, which is a hybrid of PTS+ASL. FlexRank's direct update of shared weights with mandatory nesting is theoretically superior and shows advantages at low budgets.
• vs SVD-LLM / DRONE: These are single-ratio compression methods; FlexRank produces a family of models with one training run.
• vs MatFormer / Once-For-All: These focus on width/depth/architecture; FlexRank is the first to establish elasticity in the factorization space with theoretical backing.
• vs LLM-Pruner / LayerSkip: FlexRank outperforms structured pruning and early-exit methods on Llama-3.2-1B while using significantly fewer training tokens (5B vs 839B).

Rating

• Novelty: ⭐⭐⭐⭐⭐ Elevates "nested submodel training" to a theoretical proof and combines it with DP and GAR.
• Experimental Thoroughness: ⭐⭐⭐⭐ Covers GPT-2 to Llama 3.1-8B and ViT-7B; downstream LoRA validation is a plus.
• Writing Quality: ⭐⭐⭐⭐⭐ Motivation is clearly articulated through three theorems; logical flow is exemplary.
• Value: ⭐⭐⭐⭐⭐ "Train-once, deploy-everywhere" addresses a real pain point in heterogeneous deployment.

Related Papers

• [ICLR 2026] Implicit Bias and Loss of Plasticity in Matrix Completion: Depth Promotes Low-Rank
• [ACL 2026] KoCo: Conditioning Language Model Pre-training on Knowledge Coordinates
• [NeurIPS 2025] Breaking the Frozen Subspace: Importance Sampling for Low-Rank Optimization in LLM Pretraining
• [ICML 2026] Towards Understanding Continual Factual Knowledge Acquisition of Language Models: From Theory to Algorithm
• [ICML 2026] Model Merging Scaling Laws in Large Language Models