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

Azal Ahmad Khan; Keshav Ramji; Tahira Naseem; Ali Anwar; Ramón Fernandez Astudillo

cs.LGOct 1, 2026

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

Authors: Azal Ahmad Khan, Keshav Ramji, Tahira Naseem, Ali Anwar, Ramón Fernandez Astudillo

Organizations: University of Minnesota · IBM Research AI

Abstract

Hypernetworks have recently shown success in dynamically adapting the parameters of Large Language Models (LLMs) at runtime based on signals such as task descriptions or additional demostrations. Here we ask: how much adaptation signal can be obtained using only the input query to an LLM?. To answer this, we study query-conditioned Hypernetworks for LoRA estimation. Further, we introduce distributional Hypernetworks, able to produce not only point estimates of parameter adaptors, but also a distribution over possible LoRAs. For this we propose a simple end-to-end loss using a differentiable Monte Carlo approximation and explore multiple distribution parametrizations including regression and convex combination variants. Results show that even using the mean of the learned distribution can outperform deterministic hypernetworks. Crucially, the learned distribution enables a different form of test-time scaling: instead of spending additional compute only by sampling more token sequences from a fixed model, we sample weight updates, yielding multiple adapted models for the same query. Performance improves as more weight samples are considered and remains stronger than corresponding token-sampling adaptation baselines. Finally, we find that generated updates can transfer across queries, suggesting that the hypernetwork learns reusable structure in how the model should adapt. Together, these results show that query-conditioned distributions over weight updates can support both adaptation and test-time scaling.

Figures & tables

Figure 1: Distributional hypernetworks predict a query-conditioned distribution over LoRA updates. The regression hypernetworks model the distribution directly in LoRA weight space, while the mixing hypernetwork models a distribution over combinations of reference LoRAs.

Panel of Figure 2 — Figure 2: Design choices for learning the weight distribution. (a) Diagonal covariance performs comparably to more expressive covariance parameterizations. (b) Best-of- $K$ training preserves a more searchable distribution as the number of sampled updates increases.

Figure 3: (a) Effect of LoRA rank on adaptation performance. Increasing the rank improves the Hypernetwork performance over FullFT. (b) Macro accuracy over seven benchmarks across five base models. Distributional training generally improves the predicted update.

Figure 4: Scaling test-time compute by sampling weights. We report macro accuracy over seven benchmarks as the sampling budget increases. For each query, we generate

K

candidate answers and use majority voting to obtain the final prediction. Dashed lines correspond to token sampling from a fixed model, including baselines and deterministic hypernetworks. Solid lines correspond to distributional hypernetworks, where each candidate is generated using a sampled weight update.

Figure 5: Generated update distributions transfer across queries. We aggregate hypernetwork outputs from one half of each benchmark and apply the resulting adaptation to the held-out half.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 6: Hypernetwork architecture. Each layer–module pair attends over the token-level query representation and produces module-specific LoRA factor parameters.

Figure 7: Training the hypernetwork. The hypernetwork generates an update

\Delta W_{x}

, which is added to the frozen language model weights

W_{0}

. The language modeling loss is backpropagated through the adapted model and the generated update to train only the hypernetwork parameters.

			Regression Hypernet		Mixing Hypernet + Bank
Base model	FullFT	LoRA $r=8$	Det.	Dist.	Det.	Dist.
Qwen3-1.7B	1,700M	5.5M	194M	295M	368.0M	368.1M
Qwen3-4B	4,000M	10.6M	245M	396M	624.1M	624.2M
Gemma-3-1B	1,000M	5.0M	192M	291M	343.3M	343.4M
Llama-3.2-1B	1,000M	3.9M	219M	345M	289.3M	289.4M
Llama-3.2-3B	3,000M	7.6M	233M	372M	472.4M	472.5M

Appendix

Table 1: Parameter counts across base models and adaptation methods. FullFT reports the number of trainable LM parameters, while LoRA reports the trainable parameters of a single rank-

8

adapter. For mixing hypernetworks, the reported count includes both the hypernetwork and the fixed bank of 50 rank-

8

reference LoRAs.

Qwen3-1.7B
Approach	Rank 2	Rank 8	Rank 32	Rank 128
FullFT trainable	$\sim$ 1,700M	$\sim$ 1,700M	$\sim$ 1,700M	$\sim$ 1,700M
LoRA trainable / generated update	1.4M	5.5M	22.0M	88.1M
Deterministic hypernet trainable	118M	194M	497M	1,707M
Qwen3-4B
Approach	Rank 2	Rank 8	Rank 32	Rank 128

Appendix

Table 2: Parameter scaling with LoRA rank for Qwen3-1.7B and Qwen3-4B. The LoRA row gives the number of trainable parameters for standard LoRA, which is also the size of the update generated by the hypernetwork.

Figure 8: Effect of LoRA placement. Hypernetwork-generated updates follow the same trend as directly optimized LoRAs: adapting the MLP is more effective than adapting attention alone, while adapting both provides only a small additional improvement.

Base model	Hypernet size	$K=1$	$K=8$	$K=64$
Qwen3-1.7B	25M	44.7	51.9	53.3
	100M	45.3	53.0	53.7
	500M	44.5	53.2	53.9
	1B	46.7	53.5	54.2
Qwen3-4B	25M	51.3	60.1	61.2
	100M	50.1	60.5	61.3

Appendix

Table 3: Effect of hypernetwork capacity on distributional weight sampling. We report accuracy for different hypernetwork sizes and sampling budgets. Increasing hypernetwork capacity generally improves single-sample performance, while the gap between small and large hypernetworks becomes much smaller as the number of sampled weight updates increases.

Figure 9: The benefit of generating updates grows with task diversity. We report the percentage of domains where the hypernetwork outperforms each baseline. Higher values indicate a larger advantage for the hypernetwork.

	Qwen3 1.7B				Qwen3 4B				Gemma 3 1B				Llama 3.2 1B				Llama 3.2 3B
	Regression		Mixing		Regression		Mixing		Regression		Mixing		Regression		Mixing		Regression		Mixing
Benchmark	Det.	Dist.	Det.	Dist.	Det.	Dist.	Det.	Dist.	Det.	Dist.	Det.	Dist.	Det.	Dist.	Det.	Dist.	Det.	Dist.	Det.	Dist.
GSM8K	56.8	64.0	78.0	69.2	64.4	62.4	88.0	82.4	29.2	28.8	30.4	41.6	44.0	40.8	49.2	52.0	68.8	72.4	82.8	84.4
MATH500	34.4	42.0	46.0	47.2	52.8	56.0	56.4	45.2	20.0	21.2	18.4	26.4	22.0	23.6	24.8	23.2	39.2	38.0	40.4	37.6
HumanEval	45.1	47.6	48.2	59.2	59.2	61.0	66.5	65.2	7.9	12.8	15.8	24.4	30.5	31.7	36.0	38.4	45.7	48.8	48.8	51.2
ARC-C	59.2	68.0	71.6	76.0	86.0	85.2	83.6	86.4	23.6	30.4	22.0	43.6	19.6	27.2	50.0	54.8	65.2	72.8	76.0	74.4

Appendix

Table 4: Benchmark-level comparison of deterministic and distributional training for regression and mixing hypernetworks across five base models on the seven benchmarks used in our main evaluation. Accuracy is reported in percentage points. Higher is better.

Model	Method	Samp.	GSM8K	MATH	HumanEval	ARC-C	MMLU	MedQA	GPQA	Macro
Qwen3 1.7B	Full FT	T	.508/.700/.768	.316/.444/.540	.529/.643/.714	.608/.740/.788	.496/.664/.700	.272/.424/.428	.217/.273/.354	.421/.555/.613
	LoRA	T	.568/.620/.632	.316/.428/.500	.667/.607/.635	.520/.760/.800	.492/.664/.708	.172/.336/.372	.177/.283/.268	.416/.528/.559
	MoL	T	.468/.560/.608	.292/.428/.480	.707/.701/.739	.596/.704/.740	.512/.688/.684	.112/.320/.352	.202/.293/.263	.413/.528/.552
	Reg. Det.	T	.524/.588/.612	.312/.420/.504	.633/.702/.738	.528/.736/.748	.496/.676/.688	.196/.356/.368	.207/.303/.308	.414/.540/.567
	Reg. Dist.	W	.620/.684/.740	.396/.528/.556	.688/.705/.706	.708/.760/.784	.608/.692/.720	.308/.444/.428	.197/.278/.258	.504/.584/.599
	Mix. Det.	T	.776/.844/.864	.456/.528/.552	.739/.781/.791	.648/.748/.752	.608/.744/.760	.280/.388/.416	.212/.298/.293	.531/.619/.633

Appendix

Table 5: Benchmark-level test-time scaling summary. Each cell reports accuracy at

K=1/8/32

\mathrm{T}

denotes token sampling from fixed adapted weights and

\mathrm{W}

denotes sampling a new weight update for each generation. Regression and mixing hypernetwork should be interpreted separately because their training exposure is not matched.

Explore similar work

Scaling Laws for Hypernetwork-Based Knowledge Injection in Large Language Models
Jul 21, 2026Nischay Dhankhar, Dos Baha, Abulhair SaparovHypernetwork Scaling Laws
Nace AI · Purdue University
SHINE: A Scalable In-Context Hypernetwork for Mapping Context to LoRA in a Single Pass
Feb 6, 2026Yewei Liu, Xiyuan Wang, Yansheng Mao +3Large Language Model Adaptation Low-Rank Adaptation Adapters
Institute for Artificial Intelligence, Peking University · School of Electronics Engineering and Computer Science, Peking University · University of Oxford +1
Query-Conditioned Test-Time Self-Training for Large Language Models
May 13, 2026Chaehee Song, Minseok Seo, Yeeun Seong +2Test-Time Training Test-Time Scaling
School of Electrical Engineering, KAIST · 2Graduate School of Green Growth and Sustainability, KAIST

Jul 21, 2026, cs.CLJ/K move · Enter open · S save

Scaling Laws for Hypernetwork-Based Knowledge Injection in Large Language Models

Nischay Dhankhar, Dos Baha, Abulhair Saparov

Nace AI · Purdue University

Hypernetwork Scaling Laws Multi-Hop Reasoning

Open paper AskPDF arXiv

Injecting factual knowledge into large language models (LLMs) reliably and at scale remains an open challenge. Hypernetworks provide a promising solution to large-scale knowledge injection. Although hypernetworks are typically applied for test-time adaptation, we explore their use in train-time knowledge injection, where, given a large corpus of facts, we train a hypernetwork to generate a fixed LoRA adapter that, when inserted into the target model, enable the model to answer questions about those facts. In this work, we investigate whether hypernetworks can be used to perform train-time knowledge injection and how this ability varies with scale. The scaling behavior of hypernetworks remains largely unstudied. Our design decouples the hypernetwork's injection capacity from the target model's general capability, enabling, for the first time, a rigorous study of scaling laws for hypernetwork architectures. We characterize how loss, reasoning accuracy, and out-of-distribution (OOD) generalization vary with hypernetwork depth, width, and target network size. We construct a large-scale dataset, called MegaWikiQA, containing tens of millions of multi-hop question-answer examples across 39 domains constructed from examples in Wikidata5M. Our results reveal: (i) hypernetwork-based injection exhibits broadly predictive power law scaling along all architecture axes; and (ii) hypernetworks are capable of reliable OOD generalization at increasing scales, suggesting that hypernetwork provides a promising alternative to other train-time adaptation methods such as LoRA finetuning and full fine-tuning, exhibiting steeper scaling exponents in all OOD evaluations. Together, these results establish hypernetworks as a principled and scalable substrate for train-time adaptation, and provide the first empirically grounded scaling laws to guide hypernetworks for factual reasoning in large language models.

Abstract

Explore similar work

Scaling Laws for Hypernetwork-Based Knowledge Injection in Large Language Models

SHINE: A Scalable In-Context Hypernetwork for Mapping Context to LoRA in a Single Pass

Query-Conditioned Test-Time Self-Training for Large Language Models