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

TL;DR

FlexRank addresses the rigidity of large pretrained models by learning an importance-ordered, nested low-rank decomposition that yields elastic submodels sharing weights. It initializes per-layer low-rank factors with DataSVD, then uses dynamic programming to identify nested configurations along a Pareto front, followed by distillation to refine performance without retraining from scratch. The approach achieves smoother accuracy degradation across budgets and demonstrates strong post-adaptation capabilities, enabling deployment across diverse hardware with minimal retraining. This framework significantly advances practical, cost-aware deployment of large models by enabling train-once, deploy-everywhere elasticity. $W_i = U_i V_i^ op$ and nested masks underpin the method, with distillation restoring performance after budget-driven pruning.

Abstract

The growing scale of deep neural networks, encompassing large language models (LLMs) and vision transformers (ViTs), has made training from scratch prohibitively expensive and deployment increasingly costly. These models are often used as computational monoliths with fixed cost, a rigidity that does not leverage overparametrized architectures and largely hinders adaptive deployment across different cost budgets. We argue that importance-ordered nested components can be extracted from pretrained models, and selectively activated on the available computational budget. To this end, our proposed FlexRank method leverages low-rank weight decomposition with nested, importance-based consolidation to extract submodels of increasing capabilities. Our approach enables a "train-once, deploy-everywhere" paradigm that offers a graceful trade-off between cost and performance without training from scratch for each budget - advancing practical deployment of large models.

Figures (8)

• Figure 1: FlexRank takes as input a base model, which is first decomposed by factorizing each linear layer independently. Next, a global ordering is obtained via a dynamic programming subroutine that assumes additivity of errors across layers. This global ordering is then used to extract nested submodels of different sizes, which are stochastically refined through distillation from the base model.
• Figure 2: Nested trained submodels are Pareto Elastic: Comparison of the considered submodels training strategies on the synthetic setting described in \ref{['sec:exp_controlled_details']}. Blue points visualize all $1023$ submodels, the red line represents the best models and the green one the true Pareto Front. The difference between the red and green lines is the best submodel optimality gap as per \ref{['eq:optimal_gap']}, and is zero only for NSL.
• Figure 3: FlexRank recovers the true Pareto Front in DNNs: points represent independently trained nested submodels, starting from (i) a random weights (red) or (ii) from the DataSVD (green) of a pretrained model (yellow star), with best models highlighted with dashed lines. At convergence, FlexRank recovers the (in advance unknown) Pareto front within a single set of shared weights.
• Figure 4: FlexRank has the most graceful performance degradation across parameter budget (NLP): Average downstream task accuracy over commonsense downstream datasets from lm-eval-harness
• Figure 5: FlexRank has the most graceful performance degradation across parameter budget (CV): classification accuracy on the evaluation split of ImageNet1K. The performance gap remains within a $5\%$ margin w.r.t. the full model even pruning up to $70\%$.

Theorems & Definitions (17)

• Definition 2.1: Pareto Elastic Model
• Remark 3.1
• Theorem 4.1
• Theorem 4.2: ASL has strictly positive submodel gap
• Theorem 4.3: NSL preserves nested minimizers
• Lemma 2.3: Objective equivalence without empty mask in ASL
• proof
• Lemma 2.4: Rank-dropout objective expansion
• proof
• Lemma 2.5: Balanced factorization penalty