PaperlayerView exact revision on arXiv

Source-linked AI summary

MatryoshkaLoRA: Learning Accurate Hierarchical Low-Rank Representations for LLM Fine-Tuning

Ionut-Vlad Modoranu, Mher Safaryan, Dan Alistarh

arXiv:2605.07850v1cs.CLcs.AIcs.LG

TL;DR

Fine-tuning large models remains costly, while static-rank LoRA requires rank selection and multiple evaluations. MatryoshkaLoRA trains nested low-rank prefixes using diagonal weighting, enabling multi-rank deployment and rank-performance evaluation with AURAC.

  • Problem

    Fine-tuning billion-parameter models remains computationally prohibitive, while choosing LoRA ranks requires evaluating different rank settings.

  • Method

    MatryoshkaLoRA inserts a shared diagonal weighting between LoRA adapters to learn hierarchical low-rank prefixes within the same adapters.

  • Results

    The framework supports evaluating or deploying one adapter at multiple ranks and introduces AURAC to summarize rank-performance trade-offs.

  • Takeaways & Limitations

    Hierarchical adapters can reduce repeated rank-specific training, support deployment at targeted ranks, and enable dynamic rank selection under varying cluster loads.

  • Takeaways & Limitations

    Evaluations used the same rank for the entire network, without testing layer-specific rank assignments in a forward pass.

Abstract

from arXiv · show

With the rise in scale for deep learning models to billions of parameters, the computational cost of fine-tuning remains a significant barrier to deployment. While Low-Rank Adaptation (LoRA) has become the standard for parameter-efficient fine-tuning, the need to set a predefined, static rank $r$ requires exhaustive grid searches to balance efficiency and performance. Existing rank-adaptive solutions such as DyLoRA mitigate this by sampling ranks during the training from a predefined distribution. However, they often yield sub-optimal results at higher ranks due to lack of consistent gradient signals across the full hierarchy of ranks, thus making these methods data-inefficient. In this paper, we propose MatryoshkaLoRA, a general, Matryoshka-inspired training framework for LoRA that learns accurate hierarchical low-rank representations by inserting a fixed, carefully crafted diagonal matrix $P$ between the existing LoRA adapters to scale their sub-ranks accordingly. By introducing this simple modification, our general framework recovers LoRA and DyLoRA only by changing $P$ and ensures all sub-ranks embed the available gradient information efficiently. Our MatryoshkaLoRA supports dynamic rank selection with minimal degradation in accuracy. We further propose Area Under the Rank Accuracy Curve (AURAC), a metric that consistently evaluates the performance of hierarchical low-rank adapters. Our results demonstrate that MatryoshkaLoRA learns more accurate hierarchical low-rank representations than prior rank-adaptive approaches and achieves superior accuracy-performance trade-offs across ranks on the evaluated datasets. Our code is available at https://github.com/IST-DASLab/MatryoshkaLoRA.

1 Introduction & Related Work

The paper targets the cost and rank-selection burden of LoRA fine-tuning by training nested low-rank adapters whose prefixes remain accurate across ranks. MatryoshkaLoRA unifies related parameterizations and supports evaluation or deployment at multiple ranks from one checkpoint.

  • Fine-tuning billion-parameter models remains computationally prohibitive for many applications despite their use as downstream knowledge bases.
  • LoRA reduces fine-tuning overhead but requires a predefined rank, making performance-efficiency trade-offs dependent on exhaustive searches across training runs.
  • Adaptive-rank methods improve parameter efficiency but typically produce one specialized configuration rather than a nested adapter family usable at inference.
  • DyLoRA samples one rank per step, leaving ranks above the sampled k without gradient signal and potentially failing to learn a true hierarchy.
  • MatryoshkaLoRA trains all specified rank prefixes as nested representations, enabling multiple deployment ranks from a single checkpoint without retraining separate adapters.
  • AURAC summarizes rank-performance trade-offs across evaluated ranks while weighting each rank according to its magnitude.

2 Method

The method represents LoRA rank prefixes through shared adapter factors and a fixed diagonal weighting vector P. This formulation aggregates all training-rank contributions efficiently while recovering standard LoRA and DyLoRA through alternative choices of P.

  • 2.1 Notation: LoRA uses adapters A and B with maximum rank R, while rank-k prefixes select the first k columns of A and first k rows of B.
  • 2.1 Notation: The method restricts training and inference ranks to powers of two, matching common LoRA rank settings.
  • 2.3 MATRYOSHKALORA: MatryoshkaLoRA includes every specified prefix A_r and B_r in the forward pass so lower-dimensional representations are contained within higher-dimensional prefixes.
  • 2.3 MATRYOSHKALORA: A direct masked implementation would require |S| matrix multiplications per layer and storing two masks per linear layer, creating undesired overhead.
  • 2.3 MATRYOSHKALORA: The method replaces these masks with a diagonal matrix diag(P), equivalently scaling rows of A by the R-dimensional vector P, whose single shared copy adds only R elements of memory.
  • 2.3 MATRYOSHKALORA: During training P weights the contributions of all ranks in S, whereas evaluation discards P and selects a rank k using the standard LoRA forward form.
  • 2.3 MATRYOSHKALORA: The framework recovers LoRA with P = 1/r · 1_R and DyLoRA with P = 1/k on the first k components and zeros thereafter.

2.4 Gradient Computation for MATRYOSHKALORA

MatryoshkaLoRA uses the diagonal weighting vector to scale adapter gradients so every trained rank contributes to learning hierarchical representations. AURAC aggregates accuracy across ranks while accounting for their spacing.

  • Gradient computation: The diagonal matrix diag(P) scales the optimizer gradients for adapters A and B, producing the targeted hierarchical low-rank representations.
  • Rank evaluation: AURAC evaluates a model at ranks in S by collecting the corresponding accuracies and aggregating them with the trapezoidal rule.
  • Rank evaluation: The standard AURAC weights intervals according to rank distance, so the 8-to-16 interval contributes 8/15 = 53.3% in the example.
  • Rank evaluation: Log-AURAC applies only to power-of-two ranks and weights their equally spaced log intervals uniformly; the experiments report no significant difference from AURAC.

3 Experimental Results

Across the evaluated Llama models and benchmarks, MatryoshkaLoRA delivers stronger and more consistent rank-wise accuracy than LoRA and DyLoRA, with equal running-time and memory requirements to those baselines.

  • Experimental setup: Experiments use Llama-3.2-1B-Instruct and Llama-3.1-8B-Instruct across GSM-8k, ARC-Challenge, and HellaSwag, with ranks formed from powers of two.Runs average three seeds and select the largest average AURAC; experiments use single H100 GPUs.
  • Efficiency: MATRYOSHKALORA has identical running-time and memory requirements to LoRA and DyLoRA because it adds only a column-wise multiplication by vector P.The added operation is the only stated resource overhead.
  • GSM-8k results: On GSM-8k with Llama-3.2-1B-Instruct, MATRYOSHKALORA reaches up to 39% at rank 16 and 38% at rank 32, versus approximately 35% for LoRA and 35.5% for DyLoRA.Its rank-wise accuracy increases substantially from R = 2, while both baselines show modest gains and plateau in the mid-34% range.
  • GSM-8k results: MATRYOSHKALORA reaches an AURAC of 38.4% on the Llama-3.2-1B-Instruct GSM-8k experiment, compared with peak values of approximately 35% for LoRA and DyLoRA.AURAC aggregates performance across ranks and captures the rank-performance tradeoff.
  • Open LLM Leaderboard results: On ARC-Challenge and HellaSwag, MATRYOSHKALORA exceeds 58% per-rank accuracy on ARC-C and peaks at 62.8% at rank 256 on HellaSwag.HellaSwag AURAC reaches 61.4%, while LoRA and DyLoRA remain effectively flat near the pretrained accuracy.

4 Conclusion, Limitations and Broader Impact

MATRYOSHKALORA modifies LoRA with a diagonal matrix to learn hierarchical low-rank features within shared adapters and supports rank evaluation through AURAC. The paper reports experiments on Llama models and notes limitations in evaluation cost and layer-wise rank variation.

  • Conclusion: MATRYOSHKALORA inserts a carefully crafted diagonal matrix diag(P) between LoRA adapters to embed hierarchical low-rank features as prefixes.
  • Conclusion: The framework recovers standard LoRA and DyLoRA by changing the vector P and introduces a metric for selecting models evaluated across multiple ranks.
  • Conclusion: Experiments evaluate 1-billion- and 8-billion-parameter Llama models on several literature datasets.
  • Limitations: Assessing performance across ranks requires many model evaluations, increasing the runtime needed to evaluate the method.
  • Limitations: The evaluations use the same rank throughout the network and do not test assigning different ranks to different layers in one forward pass.
  • Broader Impact: The method is intended to reduce deployment costs, but the authors do not control its downstream applications.

A Simplification of MATRYOSHKALORA

The simplification expresses MATRYOSHKALORA through elementwise-scaled LoRA factors and a diagonal matrix P. Its scaling assigns larger coefficients to earlier adapter components so all ranks contribute to hierarchical feature learning.

  • For 3×3 adapters supporting ranks 1, 2, and 3, the construction uses prefixes of A's columns and B's rows for each rank.
  • The product of elementwise-scaled factors is designed to equal the sum of rank-truncated products across the supported ranks.
  • The unscaled expansion contains the rank-1, rank-2, and rank-3 contributions through progressively longer prefixes.
  • The scaling coefficients are 3, 2, and 1, so earlier rows of A and columns of B receive larger scaling than later components.
  • Using square roots of the coefficients for the two factors preserves the desired combined scaling for both A and B.
  • Multiplying A's columns by P produces the scaled factor used in the MATRYOSHKALORA training forward pass.

B DyLoRA vs MatryoshkaLoRA: Theory View

The theory view formulates one rank-R adapter pair as a family of truncated adapters optimized across supported ranks. A weighted multi-rank objective assigns training emphasis to each rank, while standard LoRA remains tied to one rank.

  • Standard LoRA uses trainable adapters A and B with maximum rank R to modify a pretrained linear layer.
  • The rank-dependent scaling factor s_r controls the scaling used for each truncated adapter.
  • A single rank-R adapter generally requires retraining when a smaller rank is desired because naive truncation is not explicitly trained for performance.
  • The goal is to train one pair of rank-R adapters whose truncated prefixes provide useful adaptations for every supported rank r in S.
  • The multi-rank objective combines rank-specific losses, with weights λ_r controlling the training emphasis assigned to each rank.

B.1 DyLoRA as stochastic optimization of the multi-rank objective

DyLoRA optimizes the multi-rank objective stochastically by sampling one rank per training step and using the corresponding truncated adapter. Its sampled gradient is an unbiased estimator of the full multi-rank gradient.

  • DyLoRA samples a rank b from a predefined distribution over a contiguous rank range at each training step.
  • The forward pass uses the truncated perturbation Δ_b = s_bA_bB_b with the sampled rank's scaling factor.
  • Each training step computes the loss for the network using only the sampled rank-b adapters.
  • DyLoRA corresponds to the weighted multi-rank objective when λ_r equals the sampling probability p_B(r).
  • The sampled DyLoRA gradient is an unbiased estimator of the full multi-rank gradient.
  • Sampling one rank-loss term avoids computing the full sum over ranks at every step while optimizing the same expected objective.

B.2 From DyLoRA’s stochastic objective to a MatryoshkaLoRA surrogate

The derivation connects DyLoRA’s sampled-rank objective to a deterministic MatryoshkaLoRA surrogate. Instead of sampling one rank per step, the surrogate combines all nested rank components in one forward pass.

  • Objective construction: The derivation starts by defining the loss through the LoRA perturbation and representing each rank component as Δ_r = s_rA_rB_r = s_rAP_rB.The rank-indexed perturbations provide the components used in the multi-rank objective.
  • Objective construction: Under differentiability and L-smoothness, first-order expansions relate losses at different nested perturbations.The smoothness condition bounds gradient variation by the Frobenius distance between perturbations.
  • Surrogate derivation: Convexity of the squared norm and subtraction of the resulting identities yield a deterministic surrogate objective.These steps transform the multi-rank formulation into an objective suitable for a single forward computation.
  • Surrogate derivation: Instead of sampling one rank per step as in DyLoRA, MatryoshkaLoRA uses one forward pass with a weighted combination of all nested rank components.The weighting is implemented through an intermediate diagonal matrix P.

B.3 Summary of the connection

The connection interprets MatryoshkaLoRA as a deterministic first-order surrogate for DyLoRA’s sampled-rank training. Under smoothness and small-perturbation assumptions, the method uses diagonal weighting to train nested rank components jointly, though the reduction is only local.

  • Connection summary: DyLoRA samples one rank during training, whereas MatryoshkaLoRA replaces this stochastic procedure with a deterministic first-order surrogate.The surrogate is designed to represent the multi-rank training problem in a single objective.
  • Connection summary: The surrogate uses an intermediate diagonal weighting matrix P to combine nested rank components within one LoRA-style objective.This weighting reflects how strongly different truncation levels are emphasized during training.
  • Scope of the connection: The reduction is local rather than exact, with quadratic error in the perturbation size under the stated smoothness and small-perturbation assumptions.It frames training as learning a shared rank-R factorization whose components receive truncation-level-dependent weights.
Loading 2605.07850v1…