Learning to Predict Distributions over Weight Updates for Test-Time Adaptation

1 INTRODUCTION

• Hypernetworks predict target-network weights after offline training, replacing repeated test-time optimization with direct prediction.The introduction frames their learned parameters as slow weights that generate fast weights for each new input.Source: §1 · p. 1
• The learned distribution’s mean outperforms deterministic hypernetworks by 4.6–11.7% accuracy across 9 benchmarks.Sampling adapted models is competitive with, and often stronger than, sampling more outputs from fixed weights; generated updates can also transfer to held-out queries.Source: §1 · p. 2
• The paper studies whether a single query contains enough adaptation signal to generate useful LoRA updates without task-level auxiliary information.Earlier hypernetwork approaches condition on task descriptions or demonstrations, while this approach uses the input query.Source: §1 · pp. 1–2
• Distributional hypernetworks use a differentiable Monte Carlo approximation to predict query-conditioned distributions over plausible LoRA updates.Variants model distributions directly in LoRA weight space or over convex combinations of reference LoRAs.Source: §1 · p. 2, Figure 1 · p. 2
• The proposed test-time scaling spends additional inference compute on sampling alternative model updates rather than additional token sequences.The paper also examines when query-specific generation is useful and whether adaptations generalize to other in-domain queries.Source: §1 · p. 2

2 RELATED WORK

• Earlier hypernetwork methods predict a single adaptation, whereas this work predicts query-conditioned LoRA updates and distributions over possible updates.The cited approaches generate adapters from demonstrations or instructions, task descriptions [6], longer contexts or related examples [9] [10], or individual inputs.Source: §2 · p. 2
• Learned weight distributions have also represented uncertainty, supported ensembling, or modeled a global distribution over model parameters.Examples include Bayesian Hypernetworks [19], implicit weight uncertainty [20], and HyperGAN [21].Source: §2 · pp. 2–3
• Unlike prior low-rank adaptation distributions, this distribution is conditioned on each query and trained through the frozen language model’s response loss.ICM-LoRA [22] and DiffLoRA [23] learn from collections of optimized adapters; Bayesian LoRA methods model task-level adaptation uncertainty [24] [25], and LoRA ensembles use independently trained task-level adaptations [27].Source: §2 · p. 3

3.1 HYPERNETWORK FOR LORA ESTIMATION

• Only hypernetwork parameters Φ are optimized offline through the frozen LM; at test time, one forward pass maps the query to an update without responses, backward passes, or iterative optimization.The pretrained parameters θ remain fixed, and the adapted model generates the response.Source: §3.1 · p. 3
• LoRA factors reduce each matrix update from d_out × d_in values to r(d_in + d_out), while the hypernetwork conditions on query hidden states and layer and module embeddings.The query supplies semantic conditioning; layer and module embeddings identify where updates are applied.Source: §3.1 · pp. 3–4

3.2 DISTRIBUTIONAL HYPERNETWORK

• A single useful update is underdetermined because multiple adaptations may yield good responses, so a point estimate can average plausible solutions into an ineffective one.The method therefore predicts a distribution over query-specific adapters.Source: §3.2 · p. 4
• Inference can use the expected LoRA update or sample K adapters and greedily decode from each, with predictive diversity coming from weight samples rather than token samples.The sampled candidates can be combined using a task-specific minimum Bayes risk rule, including majority voting.Source: §3.2 · p. 4
• The Gaussian variant samples LoRA factors using a reparameterized mean-and-covariance transform, exploring diagonal or low-rank covariance to avoid full covariance’s quadratic output burden.A full covariance requires n(n+1)/2 Cholesky parameters.Source: §3.2 · pp. 4–5
• The convex-hull variant maps Gaussian mixing logits through softmax to combinations of reference adapters, restricting samples to their convex hull in a lower-dimensional region of weight space.Its logit space has R − 1 effective degrees of freedom.Source: §3.2 · p. 5

Loss reductions

• Mean loss is the unbiased Monte Carlo estimator of expected risk and gives every sampled update a gradient.The per-example gradients are averaged across mini-batches.Source: §4 · p. 5
• Best-of-K training preserves more searchable distribution as the number of sampled updates grows, while the minimum-loss objective rewards search coverage.The minimum selects the highest-likelihood sampled adapter, but its gradient depends on a single winning sample and may favor lucky samples or an early-winning mode.Source: Figure 2 · p. 6, §4 · p. 6

Covariance parameterization

• In Figure 2's design-choice comparison, diagonal covariance performs comparably to more expressive covariance parameterizations, while the caption separately reports that best-of-K training preserves a more searchable distribution as sampled updates increase.Source: Figure 2 · p. 6

Loss reductions

• Mean and especially maximum-loss training tend to shrink learned covariance without an explicit entropy reward, variance constraint, or prior.Maximum loss targets robustness but can reduce worst-case loss by collapsing the learned distribution.Source: §4 · p. 6
• Minimum, mean, and maximum reductions respectively favor search coverage, average sampled performance, and robustness.Minimum and maximum optimize K-dependent order statistics and are biased relative to expected risk.Source: §4 · pp. 5–6

5.1 EXPERIMENTAL SETUP

• Experiments use five language models, seven evaluation benchmarks, and roughly 1.3 million training examples; the architecture predicts query-conditioned updates or mixture weights over reference LoRAs.The study also explores covariance and loss-reduction variants, and compares against FullFT, LoRA, and Mixture-of-LoRA.Source: §5.1 · pp. 6–7

5.2 PREDICTING WEIGHT UPDATES FROM QUERY

• Hypernetworks generally outperform other methods, and generated LoRAs outperform full fine-tuning at rank 32 or higher while remaining smaller than full-model updates.The evaluation covers five language models and seven benchmarks; all methods use rank 8 except in the rank sensitivity analysis.Source: §5.1 · pp. 6–7, §5.2 · p. 7

5.3 PREDICTING WEIGHT DISTRIBUTION FROM QUERY

• Weight sampling remains competitive with token sampling across five models and seven benchmarks and is often stronger than the corresponding deterministic hypernetworks.Each distributional candidate uses a separately sampled weight update, whereas deterministic methods sample outputs from fixed adapted weights.Source: Figure 4 · p. 8, §5.3 · pp. 8–9

5.2 PREDICTING WEIGHT UPDATES FROM QUERY

• The mean of a learned distribution outperforms deterministic hypernetworks in all but the largest model sizes, with sampling during training proposed as a source of regularization.The objective evaluates perturbed updates around the predicted mean, favoring updates whose nearby perturbations also have low loss.Source: §5.2 · pp. 7–8

5.4 REUSING GENERATED UPDATES ACROSS QUERIES

• Aggregated updates from one half of a benchmark remain effective on the held-out half, and distributionally trained hypernetworks generally produce stronger reusable adaptations than deterministic counterparts.The hypernetwork is not called on held-out queries, and the first-half examples provide no labels or optimization.Source: §5.4 · p. 9

• A single query can provide enough adaptation signal to generate useful query-conditioned LoRA updates and distributions over plausible updates.Source: §6 · p. 9
• Sampling different adapted models is competitive with, and often stronger than, sampling more outputs from fixed weights.Source: §6 · p. 9
• Generated updates exhibit reusable structure across queries, supporting per-query and shared task-level adaptation.Source: §6 · p. 9

• Experiments cover 1–4B models with LoRA adaptation; behavior on substantially larger models remains untested.Source: §7 · p. 10
• The hypernetwork requires memory and one forward pass per query, but avoids optimization-based adaptation’s backward pass, optimizer state, and iterative updates.Source: §7 · p. 10
• The authors identify larger-model scaling and more efficient hypernetworks as promising future directions.They motivate this direction by the reported consistent gains from generated and sampled updates.Source: §7 · p. 10

• The paper provides information to reproduce the experiments across the main paper and appendix.Source: REPRODUCIBILITY STATEMENT · p. 10
• Reported reproducibility details include evaluation models and benchmarks, training datasets, and important training and implementation hyperparameters.The appendix supplies additional experimental and implementation details.Source: REPRODUCIBILITY STATEMENT · p. 10

AI USE STATEMENT

• AI tools assisted with code writing and editing and manuscript editing; the authors reviewed the assisted material and take responsibility for the final content.Source: AI USE STATEMENT · p. 10

A EXTENDED RELATED WORK

• Prior test-time training methods optimize after seeing the test input, whereas this work trains an adaptation mechanism beforehand to produce a query-conditioned update without test-time gradient optimization.Source: A EXTENDED RELATED WORK · p. 14
• Prior hypernetworks generate adaptations from demonstrations, instructions, task descriptions, longer contexts, or related examples; this work uses only the current query and predicts a distribution over LoRAs.Source: A EXTENDED RELATED WORK · p. 14
• Earlier weight-distribution methods model global weights or reproduce optimized adapters, while this distribution is query-conditioned and trained through the frozen language model’s response loss.Samples represent alternative adaptations for the current query, and the predicted distribution can vary across queries from the same task.Source: A EXTENDED RELATED WORK · p. 14
• Most test-time scaling methods sample more output sequences from fixed weights; this approach learns a query-conditioned search distribution and samples alternative adapted models.The mixing variant samples combinations of reference LoRAs rather than returning just one composition.Source: A EXTENDED RELATED WORK · pp. 14–15

B.2 TRAINING DETAILS

• Training samples candidate LoRA updates and backpropagates language-model loss through the adapted model and sampled update, updating only hypernetwork parameters.The pretrained language-model weights remain frozen; reparameterization lets gradients train both the predicted distribution’s center and spread.Source: §B.2 · p. 15, Figure 7 · p. 16

B.3 HYPERNETWORK PARAMETER COUNT

• Hypernetwork parameter counts vary with base-model dimensions and grow with LoRA rank; diagonal-Gaussian distributional versions add σ heads for the predicted spread.Table 1 compares FullFT, standard rank-8 LoRA, and hypernetwork counts; mixing counts include a fixed bank of 50 rank-8 reference LoRAs.Source: §B.3 · pp. 15–16, Table 1 · p. 17, Table 2 · p. 17

B.1 HYPERNETWORK ARCHITECTURE

• Each layer–module pair attends to query-token representations, and its output heads produce module-specific LoRA factors; distributional heads predict factor means and element-wise standard deviations.The architecture uses learned layer and module embeddings, cross-attention, residual MLP blocks, and separate heads for factors A and B.Source: §B.1 · p. 15, Figure 6 · p. 16

C IMPLEMENTATION DETAILS

• Training uses antithetic Gaussian-noise pairs, ε and −ε, to reduce variance introduced by sampling.Source: C IMPLEMENTATION DETAILS · p. 16
• A convex mixture of reference LoRA updates is applied exactly as one rank-Rr LoRA; it adds no parameters, though adapter-path compute grows with R.The reference-adapter factors are views of a frozen bank, and the exact formulation preserves the cross terms of the update mixture.Source: C IMPLEMENTATION DETAILS · p. 17

D.1 EFFECT OF HYPERNETWORK-GENERATED LORA PLACEMENT

• MLP-targeted generated LoRAs improve more than attention-only updates, while adapting both gives only a small additional gain.Generated updates follow the same qualitative placement trend as directly optimized LoRAs.Source: Figure 8 · p. 18, §D.1 · p. 18

C IMPLEMENTATION DETAILS

• Increasing hypernetwork capacity generally improves single-sample accuracy, while the gap between small and large hypernetworks shrinks as more weight updates are sampled.Source: Table 3 · p. 18

D.2 EFFECT OF HYPERNETWORK CAPACITY.

• 4.5 accuracy points at K = 1 versus 0.9 at K = 64: scaling the Qwen3-4B hypernetwork from 25M to 1B parameters helps single-update performance more.Because sampling multiple updates reduces dependence on a very large hypernetwork, the main experiments use moderately sized hypernetworks.Source: §D.2 · p. 18

E WHEN DOES GENERATING AN UPDATE BECOME USEFUL?

• As task diversity grows, query-conditioned generation gains an advantage over shared or fixed-bank updates, especially on unseen domains.On in-domain data, the hypernetwork consistently outperforms a shared LoRA; on unseen domains, the shared LoRA is competitive at low diversity but less effective as diversity rises.Source: Figure 9 · p. 19, E WHEN DOES GENERATING AN UPDATE… · p. 19

F.1 CHOOSING THE COVARIANCE

• At K = 8, diagonal covariance reaches approximately 0.75 mean accuracy, versus 0.74 for block diagonal and 0.68 for full covariance across 8 benchmarks.The three parameterizations perform similarly across sampling budgets, so the simpler, cheaper diagonal covariance is used in remaining experiments.Source: §F.1 · p. 19

F.2 LEARNING A SEARCHABLE DISTRIBUTION

• The objective controls the balance between moving the distribution’s center and preserving its spread: best-of-K pressures the full distribution least, while top-fraction objectives are intermediate.With larger K, a good sample can lie farther from the mean, making the objective more tolerant of a wider distribution.Source: §F.2 · p. 20
• Best-of-K retains substantially more coverage as the weight-sampling budget increases than objectives that require many sampled updates to perform well.Objectives placing more pressure on the full distribution leave less room for search.Source: §F.2 · p. 20

G.1 BENCHMARK-WISE COMPARISON BETWEEN DETERMINISTIC AND DISTRIBUTIONAL HYPERNETWORKS

• Across five models, distributional training generally improves macro performance over a directly trained deterministic hypernet.Table 4 covers regression and mixing hypernetworks on seven benchmarks; accuracy is reported in percentage points, with higher values better.Source: §G.1 · p. 20

G.2 FULL BENCHMARK-LEVEL TEST-TIME SCALING RESULTS

• Table 5 reports accuracy at K = 1/8/32, comparing token sampling from fixed adapted weights (T) with fresh weight-update sampling (W).It covers seven benchmarks and their macro average; regression and mixing hypernetworks should be interpreted separately because their training exposure is unmatched.Source: Table 5 · p. 21, §G.2 · p. 21