MatryoshkaLoRA: Train Once, Get Effective LoRA Adapters at Every Rank

> Paper: MatryoshkaLoRA: Learning Accurate Hierarchical Low-Rank Representations for LLM Fine-Tuning > Authors: Ionut-Vlad Modoranu (ISTA), Mher Safaryan (Lancaster University), Dan Alistarh (ISTA) > Paper: https://arxiv.org/abs/2605.07850 > Code: https://github.com/ISTA-SYNS/GLARE

---

1. LoRA's Pain Point: Choosing a Rank Is a Gamble

LoRA is the most popular method for fine-tuning large language models, but it has a frustrating problem:

The rank must be fixed in advance.

Too small, and performance suffers; too large, and parameters are wasted. Which rank is best? You don't know—so you grid-search: train 5 or 10 times at Rank=4, 8, 16, 32 and pick the winner.

That's wasteful. DyLoRA tried to improve on this by randomly sampling one rank per step, hoping a single training run covers multiple ranks. But DyLoRA has a fundamental flaw:

> Only one rank receives gradient at each step; the other ranks' columns/rows are never trained at that step.

That's not true hierarchical learning—it's a blind lucky draw.

---

2. The Core Idea: One Vector P So All Ranks Learn Simultaneously

2.1 Inspiration from Matryoshka Dolls

A Matryoshka (Russian nesting) doll has the property that: a large doll contains a smaller one, and every layer is complete.

MatryoshkaLoRA works the same way: train one adapter at a maximum rank R, but its low-rank slices (R=1, 2, 4, ...) are all valid and well-trained.

2.2 Key Design: The Diagonal Vector P

Between LoRA's A and B matrices, insert a fixed diagonal scaling—essentially a vector P:

Standard LoRA:

MatryoshkaLoRA:

where \(P \in \mathbb{R}^R\) is an R-dimensional vector that rescales the columns of A via element-wise multiplication \(A * P\).

How P is constructed (e.g., R=8, S={1,2,4,8}):

| Component of P | Value | Meaning | |---------|-----|------| | p₁ | 4 | Column 1 is used by all 4 ranks (1,2,4,8) | | p₂ | 3 | Column 2 is used by 3 ranks (2,4,8) | | p₃, p₄ | 2 | Columns 3–4 are used by 2 ranks (4,8) | | p₅–p₈ | 1 | Columns 5–8 are used only by rank 8 |

Key insight: each component \(p_r\) counts how many sub-ranks share column r. Columns serving lower ranks are shared by more ranks and therefore receive stronger gradient signals.

2.3 Why It Works: Gradient Propagation

During training, the gradients are:

• \(\nabla_A = \Delta \cdot B^\top \cdot \text{diag}(P)\)
• \(\nabla_B = \text{diag}(P) \cdot A^\top \cdot \Delta\)

Effect: the gradient of column r is amplified by a factor of \(p_r\). Rank 1 uses only column 1, but column 1 is shared by all ranks, so its gradient is amplified 4×. Rank 2 uses the first two columns, whose gradients are amplified 4× and 3× respectively.

This ensures: low-rank prefixes get stronger learning signals, and representations at all ranks are optimized simultaneously.

---

3. Comparison with DyLoRA: True Hierarchy, Not a Lucky Draw

| Dimension | DyLoRA | MatryoshkaLoRA | |:---|:---|:---| | Training mechanism | Samples one rank per step, using only the first k columns | Uses a weighted combination of all ranks at every step | | Gradient signal | Sparse: only the first k columns get gradient per step | Dense: all columns get gradient every step, weighted by P | | Data efficiency | Low: needs more steps to cover all ranks | High: learns all hierarchical representations every step | | Hierarchical property | Pseudo-hierarchy: high-rank subsets not directly trained | True hierarchy: low-dim representations nested in high-dim ones | | Nature | Monte-Carlo approximation of multi-objective random optimization | Deterministic surrogate: multi-objective via a single forward pass |

DyLoRA's theoretical weakness:

DyLoRA is equivalent to stochastic optimization:

\[\min_{A,B} \mathbb{E}_{k}[L_k(A,B)] = \min_{A,B} \sum_{r \in S} \lambda_r \ell(W_0 + s_r A P_r B)\]

with error \(O(\sum \lambda_r \|\Delta_r\|_F^2)\); when perturbations are large, the approximation breaks down.

MatryoshkaLoRA's solution: by explicitly constructing P, it converts multi-objective optimization into a single forward pass:

\[\sum_{r \in S} \lambda_r f(\Delta_r) \approx f\left(\sum_{r \in S} \lambda_r \Delta_r\right) = f(A P B)\]

The error is controllable, and all ranks share the same gradient flow.

---

4. AURAC: A New Metric for Evaluating Multi-Rank Adapters

4.1 Why a New Metric?

Traditional LoRA evaluation only cares about accuracy at a single rank. MatryoshkaLoRA requires evaluating the accuracy curve across all ranks.

4.2 AURAC Definition

Given a rank set \(S = \{r_1, r_2, ..., r_{|S|}\}\) with accuracies \(A_S = \{a^{(r)}\}\):

AURAC (linear version):

\[\text{AURAC} = \frac{1}{r_{|S|} - r_1} \sum_{i=1}^{|S|-1} \frac{a^{(r_i)} + a^{(r_{i+1})}}{2} \cdot (r_{i+1} - r_i)\]

log-AURAC (log version, when ranks are powers of 2):

\[\text{log-AURAC} = \frac{1}{\log_2(r_{|S|}) - \log_2(r_1)} \sum_{i=1}^{|S|-1} \frac{a^{(r_i)} + a^{(r_{i+1})}}{2} \cdot (\log_2 r_{i+1} - \log_2 r_i)\]

4.3 Design Rationale

• Trapezoidal rule: numerical integration reflecting the area under the rank-accuracy curve
• Rank weighting: larger rank intervals contribute more weight, consistent with "larger ranks should be stronger"
• Hierarchical quality: a high AURAC requires all intermediate ranks to stay accurate, not just the peak rank
• log-AURAC: when S={1,2,4,8}, all intervals get equal weight, eliminating scale bias

---

5. Experimental Results: One Training Run, Beating Everything

5.1 Llama-3.2-1B on GSM-8K (math reasoning)

| Method | AURAC | r=1 | r=4 | r=8 | r=16 | r=32 | |:---|:---|:---|:---|:---|:---|:---| | LoRA | 34.5% | - | - | - | - | 34.9% | | DyLoRA | 34.9% | - | - | - | - | 35.4% | | MatryoshkaLoRA | 38.4% | 35.9% | 37.0% | 38.8% | 39.1% | 38.8% |

Key findings:

• Beats baselines even at r=1: MatryoshkaLoRA's r=1 accuracy is 35.9%, exceeding LoRA/DyLoRA at any rank
• r=4–8 reaches 37–39%, far above the baselines' best of ~35%
• No more grid search: r=4 or 8 matches or exceeds other methods' best rank

5.2 Llama-3.1-8B Multitask

| Task | LoRA | DyLoRA | MatryoshkaLoRA | Gain | |:---|:---|:---|:---|:---| | GSM-8K (3-shot) | 74.3% | 74.4% | 77.4% | +3% | | GSM-8K (8-shot) | 79.1% | 79.2% | 79.7% | +0.6% | | ARC-C | 57.0% | 57.0% | 58.0% | +1% | | HellaSwag | 59.2% | 59.2% | 61.4% | +2.2% |

Key findings:

• 3-shot approaches 8-shot: 3-shot at r=32/64 reaches 77–78%, close to 8-shot's 79.7% → shorter context, lower inference cost
• HellaSwag increases monotonically with rank: r=256 hits 62.8% vs the baseline's 59.2% (+3.6%)

5.3 Scaling Factor Ablation

| Scaling \(s_k\) | Best learning rate | AURAC | Observation | |:---|:---|:---|:---| | \(1/r\) | \(9 \times 10^{-4}\) | 37.8% | Needs 4× larger learning rate | | \(1/\sqrt{r}\) | \(3 \times 10^{-4}\) | 38.7% | Needs 9× larger learning rate | | 1 | \(1 \times 10^{-4}\) | 38.4% | Smallest learning rate, best stability |

Choosing \(s_k=1\) instead of \(1/\sqrt{r}\) or \(1/r\): the latter make large-rank contributions too small, requiring compensating larger learning rates and harder tuning.

---

6. Minimal Implementation: Just One Vector

6.1 Extra Parameters

| Component | Count | Comparison | |:---|:---|:---| | Vector P | R floats (shared globally) | Independent of layer count | | vs. per-layer mask schemes | \(2 \times |S| \times L\) boolean matrices | Per-layer independent |

An R-dimensional vector vs. per-layer boolean matrices: overhead is negligible in practice.

6.2 Compute Overhead

| Operation | Cost | Notes | |:---|:---|:---| | Training forward: \(A * P\) | Element-wise \(O(m \times R)\) | Negligible vs. matmul \(O(m \times R \times n)\) | | Core matmul: \((A*P) \times B\) | \(O(m \times R \times n)\) | Identical to LoRA | | Total training overhead | ≈ LoRA = DyLoRA | Explicitly stated in the paper |

6.3 Zero Inference Cost

| Phase | Operation | |:---|:---| | Training | Uses \((A * P) \times B\) | | Inference | Drop P, slice \(A_k \times B_k\) directly (identical to standard LoRA) |

Deployment advantage: a single checkpoint supports dynamic rank switching—no retraining, no storing multiple adapters.

6.4 A Unified Framework

With different choices of P, MatryoshkaLoRA recovers:

• LoRA: \(P = \mathbf{1}\) (all-ones vector, degenerating to standard LoRA)
• DyLoRA: P under specific sampling strategies

MatryoshkaLoRA is thus a general framework, not just another standalone variant.

---

7. Practical Implications for Engineers

7.1 No More Grid Search

Before: 1. Train rank=4 → evaluate → not good enough 2. Train rank=8 → evaluate → still not enough 3. Train rank=16 → evaluate → acceptable 4. Train rank=32 → evaluate → best

5 training runs, 5 checkpoints, 5 evaluations.

Now: 1. Train once at max rank=32 2. Get one checkpoint containing valid adapters at all ranks 3. Pick the r=4, 8, 16, or 32 slice based on deployment hardware

1 training run, 1 checkpoint, dynamic deployment.

7.2 A Boon for Mobile Deployment

Deploying LLMs on phones requires compressing parameters as much as possible, but different tasks demand different precision:

• Simple tasks: r=4 suffices
• Complex tasks: r=16 or higher is needed

MatryoshkaLoRA: one checkpoint, rank chosen dynamically per task. No need to train a separate model per rank.

7.3 From "Search at Training" to "Choose at Inference"

DyLoRA's issue: sampling is random during training, but a fixed slice is chosen at inference—and the training-time sampling distribution may not match the inference-time choice.

MatryoshkaLoRA defers the choice to inference: all ranks are fully optimized during training, and the best slice is selected at inference under resource constraints.

---

8. Limitations and Future Directions

8.1 Accuracy Drop from Rank 128→256

On the 8B model, accuracy drops when rank goes from 128 to 256; the authors suspect more training epochs are needed. This suggests:

• Very high ranks still require sufficient training
• P's construction may need adjustment for very large ranks

8.2 Only LLaMA Architectures Tested

The paper validates on Llama-3.1/3.2 only; other architectures (e.g., Mistral, Qwen, DeepSeek) are untested. Whether P's construction generalizes needs further verification.

8.3 Joint Optimization with Quantization and Pruning

MatryoshkaLoRA solves rank-selection efficiency, but LLM deployment also involves quantization (INT8/INT4), pruning, and other dimensions. Jointly optimizing these is the next step.

---

9. Conclusion: LoRA's "Matryoshka Moment"

MatryoshkaLoRA's core contribution isn't added complexity—it's simplified selection:

> "No need to pick a rank—train once, and every rank is ready."

With a single R-dimensional vector P, it achieves:

• Gradient signals at all ranks simultaneously
• Low-rank prefixes nested within high-rank representations
• Training overhead identical to LoRA
• P discarded at inference—zero extra cost
• Single-checkpoint, multi-rank deployment

Behind this lies a deep mathematical insight: converting stochastic multi-objective optimization into a deterministic single forward pass, with P's construction ensuring gradient signal is distributed evenly across all levels.

For LLM fine-tuning practitioners, this means:

> The era of grid search may be ending.

---

References

• Modoranu, I.-V., Safaryan, M., & Alistarh, D. (2026). "MatryoshkaLoRA: Learning Accurate Hierarchical Low-Rank Representations for LLM Fine-Tuning." arXiv:2605.07850
• Code: https://github.com/ISTA-SYNS/GLARE
• Hu et al. (2022). LoRA: Low-Rank Adaptation of Large Language Models.
• Valipour et al. (2023). DyLoRA: Parameter-Efficient Tuning with Dynamic Search-Free Low-Rank Adaptation.