Queryable LoRA: Instruction-Regularized Routing Over Shared Low-Rank Update Atoms

Omatharv Bharat Vaidya; Connor T. Jerzak; Nhat Ho; Chandrajit Bajaj

Source-linked AI summary

Queryable LoRA: Instruction-Regularized Routing Over Shared Low-Rank Update Atoms

Omatharv Bharat Vaidya, Connor T. Jerzak, Nhat Ho, Chandrajit Bajaj

arXiv:2605.08423v1May 8, 2026cs.LGcs.CLstat.ML

TL;DR

Static layer-local LoRA updates may not adapt to input- and depth-dependent corrections. Queryable LoRA routes shared low-rank update atoms, improving held-out performance and optimization stability across regression and LLM fine-tuning tasks with comparable trainable parameters.

Problem
Standard layer-local low-rank adapters use fixed updates even when optimal corrections vary across inputs and computation depth.
Method
The method routes shared low-rank update atoms using the current representation, earlier-layer summaries, and an optional instruction-derived semantic prior.
Results
The method improves held-out performance and optimization stability across noisy nonlinear regression and LLM fine-tuning tasks while using comparable trainable parameters to standard low-rank adaptation.
Takeaways & Limitations
Queryable update memory preserves parameter-efficient adaptation while enabling input- and depth-conditioned updates, with language guidance providing an interpretable semantic prior.
Takeaways & Limitations
Routing does not improve every benchmark and increases forward-pass computational complexity compared with static baselines.

Abstract

from arXiv · show

We present a data-adaptive method for parameter-efficient fine-tuning of large neural networks. Standard low-rank adaptation methods improve efficiency by restricting each layer update to a fixed low-rank form, but this static parameterization can be too rigid when the appropriate correction depends on the input and on the evolving depth-wise computation of the network. Our approach replaces a purely layer-local adapter with a shared queryable memory of low-rank update atoms. For each block of layers, the model forms a query from the current low-rank state and a running summary of previous blocks, uses this query to retrieve a content-dependent combination of shared update components via attention, and applies the resulting routed operator within the low-rank bottleneck. In this way, the method retains the efficiency and scalability of low-rank adaptation while allowing the effective update to vary across inputs and to share reusable structure across layers. The resulting architecture provides a principled middle ground between static LoRA-style updates and fully generated parameter updates: it remains compact and parameter-efficient while supporting dynamic, context-sensitive adaptation. Further, we incorporate instruction-regularization by augmenting routing logits with a language-induced prior over update atoms, thereby biasing the selection of low-rank transformations toward semantically relevant directions without generating unconstrained parameter updates. Experiments on noisy non-linear regression tasks and LLM fine-tuning suggest that this queryable update-memory formulation can improve final test performance and training stability compared to standard low-rank adaptation, while using a comparable number of trainable parameters.

1 Introduction

The paper introduces a queryable low-rank update memory with blockwise routing, allowing example- and depth-dependent corrections while retaining LoRA’s efficient bottleneck and reusable structure. It further adds instruction-based semantic regularization, reports empirical gains with comparable trainable parameters, and provides bounded-update guarantees.

Motivation: This design addresses LoRA’s fixed, layer-local updates by allowing corrections to vary with the current example and depth-wise computational state.The motivation is that the appropriate correction may depend on the current hidden representation and information accumulated from earlier layers.
Queryable update memory: The method uses a globally shared memory of low-rank update atoms and a blockwise router to assemble example-dependent operators inside the LoRA bottleneck.Routing depends on the current low-rank representation and a running summary of earlier blocks.
Instruction regularization: Instruction regularization uses language instructions as a semantic prior to guide atom selection toward semantically relevant low-rank transformations without unconstrained parameter changes.The prior regularizes the selection of update atoms rather than generating unrestricted parameter updates.
Empirical results: Experiments report improved held-out performance and optimization stability on noisy nonlinear regression tasks and several LLM fine-tuning benchmarks with trainable parameters comparable to standard low-rank adaptation.The reported improvements span both noisy non-linear regression tasks and LLM benchmarks.
Theoretical guarantees: Theoretical analysis proves that dynamic updates remain bounded and norm-controlled because effective updates are convex mixtures of shared atoms.The paper also states that routing weights solve a principled optimization problem so language priors guide, rather than arbitrarily override, the model’s internal state.

2 Problem Setup, Notation & Related Work

This section frames fine-tuning as adapting a frozen pretrained network with efficient low-rank updates, then identifies static, layer-local adaptation as a structural limitation. It positions the proposed method as dynamically assembling context-dependent updates from shared rank-space atoms rather than generating dense parameters directly.

Problem Setup and Notation: Fine-tuning adapts a frozen pretrained network fθ0 for a downstream task using trainable low-rank matrices Aℓ and Bℓ.LoRA applies the update ∆Wℓ = α/r BℓAℓ, with bottleneck rank r limiting adapter capacity.
Limitations of Standard LoRA: LoRA’s fixed, layer-local subspace requires the same adapter for every input and fragments trainable capacity across independent layers.These constraints can prevent different examples or depth-wise stages from receiving different corrections and force useful adaptations to be relearned across layers.
Related Work: Magnitude-direction and sharing-based LoRA variants improve optimization or reduce redundant local structure, but their final adapters remain static.The cited sharing-based methods address training efficiency and cross-layer redundancy without allowing transformations to adapt jointly to the current computation.
Proposed Approach: The proposed method dynamically assembles updates from a globally shared memory of rank-space atoms using the current internal state and an attention-based summary of preceding activations.This queryable operator replaces the rigid low-rank bottleneck and uses the computational trajectory during the forward pass.

3 Approach: Instruction Queryable Memory for Data-Adaptive PEFT

The approach replaces fixed layer-specific LoRA updates with a compact, globally shared memory of rank-space atoms retrieved through state-, depth-, and optionally instruction-conditioned routing. Instruction regularization biases this retrieval without generating adapter weights, while limiting computation through blockwise routing and sparse combinations.

Queryable Update Memory: A queryable operator replaces rigid layer-local LoRA transformations, enabling input- and instruction-dependent updates without hypernetwork-scale parameter costs.The operator is routed inside the low-rank bottleneck rather than generated as a dense parameter update.
Queryable Update Memory: Shared rank-space atoms form a dynamic vocabulary that can rotate and scale adapter directions while remaining compact.The operator mixes coordinates within the rank space and is drawn from a globally shared memory bank.
Blockwise Routing and State Summarization: Layers are partitioned into continuous blocks, with one routed operator computed per block from layer priors, current block-entry state, earlier-block summaries, and optional instructions.Blockwise computation amortizes routing and encourages consistent structural adaptations across layers.
Instruction Regularization: Instruction regularization adds a language-derived atom prior to state-dependent routing logits, then forms the operator from a sparse top-k convex combination.The language strength is controlled by τlang, while the selected atoms remain fixed and reusable.
Instruction Regularization: External instructions are optional: language shapes routing through query conditioning and atom-level priors, but does not generate adapter weights.Setting τlang = 0 and λctx = 0 yields state-dependent dynamic routing, while gℓ = 0 recovers standard LoRA.

4 Empirical Evidence

Experiments across synthetic non-convex regression, language-model fine-tuning, continual learning, and inference-time analysis indicate that queryable routing improves generalization and optimization while remaining comparatively efficient. Instruction regularization strengthens held-out performance and gradient flow, with benefits varying across task families and architectures.

Synthetic regression: The queryable adapter helps most on locally heterogeneous targets such as Dropwave, while gains are smaller on smoother or highly regular landscapes.The results are not uniformly explained by non-convexity alone; static low-rank corrections may be too rigid for heterogeneous targets.
Language-model fine-tuning: In general language-model tasks, instruction-queryable routing outperforms LoRA on every benchmark and is strongest on six of seven tasks despite near-saturated training accuracy.This pattern is interpreted in the passage as improved generalization and routing stability rather than merely increased memorization capacity.
Language-model fine-tuning: Mathematics results are heterogeneous but supportive, with queryable variants improving or matching the strongest test accuracy on several reasoning benchmarks, especially GSM8K and Numina-Math.The largest gains appear on some multi-step mathematics benchmarks, where reusable depth-wise adaptation may be beneficial.
Optimization dynamics: The instruction-queryable adapter maintains stronger gradient signals across a wider range of layers than LoRA and queryable routing without instruction regularization.Figure 3 examines per-layer adapter gradient norms after post-training.
Inference-time efficiency: The optimized queryable adapter is slower than static LoRA but substantially faster than RepLoRA, HyRA, and DoRAN, with negligible adapter-side arithmetic.The comparison matches model, target-module, rank, batch-size, sequence-length, and decoding settings.

5 Theoretical Results

The theoretical results characterize instruction-regularized routing as variational retrieval with a semantic prior, while proving bounded dynamic updates and structured gradient supervision. Together, these guarantees explain how the method gains context-sensitive flexibility without abandoning norm control or learnability.

5.1 Variational routing: Theorem 5.1 characterizes routing weights as the unique solution to a utility–KL tradeoff toward the instruction-induced semantic prior.The router selects atoms useful for the current hidden state while regularizing toward semantically relevant directions.
5.1 Variational routing: Corollary 5.1.1 shows that when the instruction prior supports the state-preferred atom, instruction regularization cannot substantially reduce state utility.This formalizes how language can guide retrieval toward semantically meaningful updates while controlling the tension between state and instruction preferences.
5.2 Stability: Theorem 5.2 proves that each routed operator lies in the convex hull of the shared atom bank and satisfies the uniform bound ∥Sb(c)∥≤RC.Because updates combine bounded shared atoms, dynamic routing preserves norm control across instructions and blocks; the result also supports bounded depth summaries.
5.3 Gradient factorization: Theorem 5.3 gives an exact blockwise gradient factorization, providing every retrieved atom with a direct gradient signal composed of low-rank layer contributions.Shared atoms consequently receive supervision from many layers and can learn reusable update directions across depth, inputs, and instructions.
5.3 Gradient factorization: Corollary 5.3.1 shows that gradient descent raises an atom’s routing probability when its direction aligns better with the block gradient than the retrieved average, and lowers it otherwise.The router therefore learns by directly comparing candidate atoms, while the shared memory is trained toward useful update directions.

6 Limitations & Conclusion

Queryable memory enables low-rank adapters to retrieve input- and depth-conditioned updates while retaining parameter-efficient fine-tuning efficiency and stability. Its benefits are benchmark-dependent, and blockwise routing adds forward-pass complexity despite remaining faster than HyRA and RepLoRA.

Conclusion: Queryable memory retrieves input- and depth-conditioned rank-space updates during the forward pass.The method uses a queryable memory of rank-space atoms for dynamic adapter updates.
Conclusion: Dynamic routing preserves parameter-efficient fine-tuning efficiency and stability while improving optimization and generalization.Language guidance can additionally provide an interpretable semantic prior.
Limitations: The queryable approach does not outperform uniformly across all benchmarks.Its advantage varies by benchmark rather than applying consistently.
Limitations: Blockwise routing increases forward-pass complexity versus a static baseline but remains substantially faster than HyRA and RepLoRA.Empirical profiling supports the relative speed advantage over these comparable dynamic methods.
Limitations: Future work should refine atom retrieval in novel feature environments and reduce routing overhead.These directions address the method’s stated limitations.

A Additional Methodological Details · A.1 Attention Updates

For attention blocks, Queryable LoRA builds a shared router state from compact summaries of the query, key, and value low-rank states. The resulting routed rank-space operator is reused across all attention projections while projection-specific LoRA factors remain distinct.

A.1 Attention Updates: Each attention block represents its token activations with a token matrix H_ℓ and applies LoRA factors separately to projection types Q, K, V, and O.The frozen attention projections are paired with projection-specific low-rank factors.
A.1 Attention Updates: For Q, K, and V, the projection input is H_ℓ, and the corresponding rank-state is formed at token level before summarization.The rank-state has r-dimensional representations across the block’s T tokens.
A.1 Attention Updates: The attention router state is the average of the query, key, and value summaries.This averaged vector replaces the block-entry rank-state sentry for attention blocks.
A.1 Attention Updates: Using this shared state avoids circular dependence on the output projection while providing the router with a compact summary of low-rank attention activity.The router therefore conditions routing on Q, K, and V information rather than on the output projection’s resulting state.
A.1 Attention Updates: The routed rank-space operator S_b(c) is reused for Q, K, V, and O within each attention block.Reuse occurs at the routed-operator level, not by replacing the local projection factors.
A.1 Attention Updates: Projection-specific LoRA factors preserve each attention projection’s inductive bias while the shared memory supplies a reusable global rank-space transformation.This separates shared routed structure from local projection-specific adaptation.
A.1 Attention Updates: Algorithmic details for the attention updates are provided in Algorithms 1 and 2.These algorithms contain the procedural specification referenced by the methodological appendix.

B Additional Experimental Results

The instruction-queryable model distributes optimization more evenly across adapted layers than LoRA, as reflected by lower gradient concentration across epochs.

The instruction-queryable model has the lowest or near-lowest gradient concentration index across epochs, while LoRA remains more concentrated.The index is the ratio of maximum to mean per-layer adapter gradient norm; lower values indicate less optimization dominance by a few layers.

B.1 Adapter Performance Comparisons

In a controlled Qwen2.5-0.5B-Instruct ARC study, queryable adapters route sparse selections from shared update atoms while retaining a compact trainable budget. The comparison indicates that structured, input-dependent retrieval improves generalization more than simply increasing adapter-generator capacity.

Experimental setup: The study compares adapter designs on AllenAI ARC with a frozen Qwen2.5-0.5B-Instruct backbone, adapting attention and MLP projections using rank-8 updates for ten epochs across three seeds.All methods train only adapter or update-memory parameters under the same broad reasoning instruction.
Routing design: Queryable variants select a sparse top-4 subset from 16 shared update atoms, while instruction-queryable routing additionally encodes the task instruction as an external prior.This makes update selection input-dependent while reusing shared update components.
Parameter efficiency: QUERYABLE LORA and INSTRUCTION-QUERYABLE LORA use only 4.456M trainable parameters, adding about 1.3% overhead relative to LoRA while enabling input-dependent update selection.By contrast, Text-to-LoRA and Doc-to-LoRA use roughly 145-148M trainable parameters, and MoE-LoRA uses 18.58M.
Generalization comparison: The controlled ARC comparison finds that structured, input-dependent retrieval over reusable update atoms improves generalization more than simply increasing adapter-generator capacity.The evidence most clearly separates routed/queryable methods from high-parameter generated baselines.

B.2 Ablation Studies · B.3 Optimization Stress Test · B.4 LLM Batch Experimentation

The studies isolate how instruction priors and state queries affect routing, identify efficient sparse-memory configurations, and test optimization in a deep, narrow regression model. Across these settings, queryable adaptation is described as achieving strong accuracy or low-loss solutions without requiring larger parameter budgets.

B.2 Ablation Studies: The ablation fixes trainable parameters at 4.46M while varying the routing signal source or structure.Variants include instruction, state-query, instruction-only, and random or adversarial routing signals.
B.2 Ablation Studies: Ablation checkpoints are selected by early stopping on highest training accuracy, using lowest training loss to break ties at 1.0 accuracy.Reported accuracies correspond to the selected checkpoint, while timing columns summarize evaluation speed.
B.2 Ablation Studies: The Pareto sweep varies rank r ∈{4, 8, 16, 32}, shared update atoms M ∈{4, 8, 16, 32}, and routing sparsity k ∈{1, 2, 4, 8}.Settings with k > M are excluded, and all configurations use the same frozen Qwen2.5-0.5B-Instruct backbone with no more than 10 epochs.
B.2 Ablation Studies: Sparse top-k routing with compact or moderate memories produces the strongest Pareto points, rather than uniformly increasing rank, atom count, or activation.Most configurations cluster at low latency with modest accuracy, while a few sparse routed configurations improve accuracy without the largest parameter budgets.
B.3 Optimization Stress Test: The optimization stress test uses a GELU multilayer perceptron with 32 hidden layers of width 32 and a scalar output head.The deeper, narrower architecture is intended to make useful update directions propagate through many layers.
B.3 Optimization Stress Test: The stress test pretrains on 3,000 source samples for 300 epochs, then adapts on 400 of 1,200 fresh target examples and tests on the remaining 800.Source data are split 80%/20%, and target functions are shifted and perturbed noisy two-dimensional regression tasks.
B.3 Optimization Stress Test: Queryable adaptation reaches low-loss target solutions across nearly all stress-test functions, whereas several static or hypernetwork-style baselines plateau at high loss or become numerically unstable.Table 9 reports best target-training MSE, while Table 10 reports best held-out target-test MSE checkpoints.

B.5 Datasets

The paper uses a broad set of post-training and evaluation datasets spanning general-task, coding, knowledge, reading-comprehension, commonsense, and mathematics-reasoning benchmarks. These datasets correspond to the general-task benchmarks in Table 10 and mathematics-reasoning benchmarks in Table 3.

Dataset suite: The dataset suite includes GPQA-Diamond, MBPP, AI2 ARC, SuperGLUE, OpenBookQA, RACE, HellaSwag, GSM8K, MATH, Orca-Math, and Numina-Math.These datasets are used for post-training and evaluation.
Benchmark organization: The listed datasets cover general-task benchmarks and mathematics-reasoning benchmarks.The general-task benchmarks appear in Table 10, while the mathematics-reasoning benchmarks appear in Table 3.

C Continual Routing Analysis · D Additional Theoretical Analysis · D.1 Setup

Continual-learning analyses show that queryable routing preserves sparse reusable memory while adapting task-specific routes, with MBPP retaining stable routing and GPQA-Diamond using a more distinct atom mixture. The theoretical setup formalizes state-, depth-, and instruction-dependent routing over shared rank-space atoms under boundedness, normalization, and fixed-active-set assumptions.

C Continual Routing Analysis: The continual-learning experiment trains sequentially on MBPP, GSM8K, and GPQA-Diamond without resetting adapters, routers, gates, keys, or shared atoms while keeping the backbone frozen.After each stage, all three tasks are evaluated and block-averaged atom-routing distributions are recorded.
C Continual Routing Analysis: Atom usage is sparse and non-uniform: a small subset of shared atoms is repeatedly reused without uniform spreading or collapse to a single atom.This pattern is consistent with queryable-memory behavior and shared low-rank corrections.
C Continual Routing Analysis: MBPP retains a stable high-mass route pattern across later stages, whereas GSM8K and GPQA-Diamond undergo localized changes in a few atoms as later tasks are learned.The observed behavior is intermediate between rigid routing and overwriting previous routes.
C Continual Routing Analysis: Final-stage atom specialization mixes broadly reusable and task-selective operators, with several near-maximal-entropy atoms and a smaller number of low-entropy atoms.High entropy indicates broad reuse across tasks, while low entropy indicates concentration on fewer tasks.
C Continual Routing Analysis: MBPP and GSM8K have relatively close final atom-usage distributions, while GPQA-Diamond is more distinct; the largest stagewise route drift occurs when GSM8K is introduced.The KL matrix captures final task separation, and the drift matrix captures route changes after new-task learning.
D.1 Setup: For each block, the router produces probabilities over shared rank-space atoms using a state-dependent query, an attention-style depth summary, and an instruction embedding.The instruction encoder maps natural-language instruction c to a fixed embedding e(c) ∈ R^dc.
D.1 Setup: Routing combines depth-summary, language-prior, and state logits, supports full or top-k active sets, and treats the active set as fixed for local differentiability statements.For sparse routing, inactive atoms receive zero routed weight; matrix inner products are Frobenius inner products.
D.1 Setup: The theoretical assumptions bound atom, adapter, state, and instruction norms, require a positive lower bound for RMS normalization, and restrict top-k results to regions with a constant active set.For top-k routing, this is ensured by a positive margin between the kth and (k + 1)st largest fused logits.

D.2 Results · D.3 Compute Resources · D.4 Dataset Details

The results establish instruction-guided routing as bounded, stable retrieval over shared low-rank atoms, with limiting regimes connecting it to state-dependent routing, static LoRA, and language-prior selection. The appendices also document approximately 720 A100-equivalent compute hours and the synthetic, pretrained-model, and public benchmark resources used.

D.2 Results: Instruction guidance changes the retrieval distribution over shared atoms while preserving a bounded retrieval-and-composition mechanism rather than generating arbitrary text-conditioned weights.The router matches hidden states while incorporating semantic preferences from the instruction.
D.2 Results: Instruction-guided retrieval loses only a small amount of state-matching score when the tempered language prior sufficiently weights the atom preferred by the hidden state.The state-prior tradeoff formalizes the tension between state matching and instruction preference.
D.2 Results: Depth summaries remain bounded and stable, and routed operators lie in the convex hull of the shared atom bank, supporting controlled use of computational history across layers.The same boundedness argument also applies to the top-k router because its weights lie in the simplex.
D.2 Results: The shared routed operator receives blockwise learning signals accumulated across adapted layers, with atom updates weighted by routing and aligned with the block gradient.The exact factorization makes the shared memory trainable through structured supervision across depth.
D.3 Compute Resources: Approximately 720 A100-equivalent hours were used across TACC, Lambda, RunPod, and local NVIDIA RTX 4090 resources.The experiments used a heterogeneous computing landscape.
D.4 Dataset Details: The study generated noisy non-convex regression samples from analytic benchmark functions and used frozen pretrained checkpoints with parameter-efficient adapters rather than redistributing original model weights.Listed backbones include Qwen2.5-0.5B, Mistral-7B-v0.1, SmolLM2-360M-Instruct, and ReasonLite-0.6B-Turbo.
D.4 Dataset Details: LLM evaluation used standard public benchmark versions, including GSM8K, MATH, Orca-Math, Omni-MATH, NuminaMath-CoT, GPQA-Diamond, MBPP, AI2 ARC, and SuperGLUE.The paper reports aggregate results and does not reproduce held-out examples or redistribute modified dataset contents.