Hyperparameter Trajectory Inference with Conditional Lagrangian Optimal Transport

Harry Amad; Mihaela van der Schaar

cs.LGMar 2, 2026

Hyperparameter Trajectory Inference with Conditional Lagrangian Optimal Transport

Authors: Harry Amad, Mihaela van der Schaar

Organizations: Department of Applied Mathematics and Theoretical Physics University of Cambridge, UK

Abstract

Neural networks (NNs) often have critical behavioural trade-offs that are set at design time with hyperparameters-such as reward weights in reinforcement learning or quantile targets in regression. Post-deployment, however, user preferences can evolve, making initial settings undesirable, necessitating potentially expensive retraining. To circumvent this, we introduce the task of Hyperparameter Trajectory Inference (HTI): to learn, from observed data, how a NN's conditional output distribution changes with its hyperparameters, and construct a surrogate model that approximates the NN at unobserved hyperparameter settings. HTI requires extending existing trajectory inference approaches to incorporate conditions, exacerbating the challenge of ensuring inferred paths are feasible. We propose an approach based on conditional Lagrangian optimal transport, jointly learning the Lagrangian function governing hyperparameter-induced dynamics along with the associated optimal transport maps and geodesics between observed marginals, which form the surrogate model. We incorporate inductive biases based on the manifold hypothesis and least-action principles into the learned Lagrangian, improving surrogate model feasibility. We empirically demonstrate that our approach reconstructs NN outputs across various hyperparameter spectra better than other alternatives.

Figures & tables

Panel of Figure 1 — Figure 1: Dots represent true samples from the temporal process across $t\in[0,1]$ , lines represent model estimated trajectories from $t=0$ to $t=1$ . Each condition has a distinct colour.

Method	NLL ( $\downarrow$ )	CD ( $\downarrow$ )
$\mathcal{K}_{I}$	$105.713$ $(2.42)$	$0.323$ $(0.003)$
$\mathcal{K}_{\theta}$	$23.008$ $(4.62)$	$0.158$ $(0.009)$
$\mathcal{K}_{I}-\hat{\mathcal{U}}$	$-0.532$ $(0.057)$	$0.016$ $(0.001)$
$\mathcal{K}_{\theta}-\hat{\mathcal{U}}$	$-0.662$ $(0.046)$	$0.016$ $(0.001)$

Table 1: NLL and CD at

t\in\{0.25,0.75\}

Method	Reward ( $\uparrow$ )
Direct	$-6.711$ $(0.070)$
MFM	$-6.561$ $(0.053)$
$\mathcal{K}_{I}-\hat{\mathcal{U}}$	$-6.397$ $(0.031)$
$\mathcal{K}_{I}$	$-6.307$ $(0.041)$
CFM	$-6.251$ $(0.028)$
NLOT	$-6.173$ $(0.038)$

Table 3: Average surrogate Reacher rewards across

\lambda_{\text{c}}\in\{2,3,4\}

Method	Reward ( $\uparrow$ )
$\mathcal{K}_{I}$	$42.84$ $(5.86)$
Direct	$49.50$ $(17.90)$
NLOT	$52.63$ $(6.52)$
$\mathcal{K}_{\theta}$	$54.44$ $(7.66)$
CFM	$69.70$ $(7.73)$
MFM	$78.44$ $(4.47)$

Table 4: Average surrogate Cancer_nl reward across

\lambda_{\text{nk}}\in\{1,2,3,4,6,7,8,9\}

Method	WD ( $\downarrow$ )
$\mathcal{K}_{I}$	$0.570$ $(0.020)$
MFM	$0.480$ $(0.020)$
CFM	$0.464$ $(0.021)$
Direct	$0.450$ $(0.013)$
$\mathcal{K}_{\theta}$	$0.404$ $(0.018)$
NLOT	$0.291$ $(0.022)$

Table 6: WD of surrogate two moons distributions from ground-truth diffusion models, trained at dropout

p\in\{0.1,0.2,0.3,0.4,0.6,0.7,0.8,0.9\}

	Semicircle		Cancer	Reacher	Dropout
$G_{\theta_{G}}$	NLL ( $\downarrow$ )	CD ( $\downarrow$ )	Reward ( $\uparrow$ )	Reward ( $\uparrow$ )	WD ( $\downarrow$ )
$R_{\theta_{G}}ER_{\theta_{G}}^{T}$	$-0.602$ $(0.033)$	$0.016$ $(0.001)$	$98.72$ $(6.32)$	$-6.122$ $(0.080)$	$0.076$ $(0.003)$
$R_{\theta_{G}}E_{\theta_{G}}R_{\theta_{G}}^{T}$	$-0.662$ $(0.046)$	$0.016$ $(0.001)$	$102.49$ $(5.46)$	$-6.093$ $(0.036)$	$0.079$ $(0.003)$

Table 7:

G_{\theta_{G}}

ablations in 2D experiments.

Appendix figures & tables10 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: Example inference-time adjustment enabled by HTI. We illustrate disparate user preferences affecting desired NN behaviour (desired $\lambda$ level) for different users in this abstract example. Having a fixed number of trained NNs ( $p_{\theta_{\lambda_{i}}}$ ) only allows partial exploration of the full hyperparameter trajectory, while an HTI surrogate model ( $(\hat{p}(\cdot|x_{i},\lambda))_{\lambda\in\Lambda}$ ) can estimate outputs across the entire spectrum of hyperparameters (estimated conditional probability paths represented by solid blue/red lines). Crucially, hyperparameter-induced dynamics can differ amongst input conditions ( $x_{i}$ ), as the true conditional distributions move along their own respective manifolds ( $\mathcal{M}_{x_{i}}$ ), so an effective HTI model must learn conditional dynamics.

Hyperparameter	Value
$\alpha$	$0.05$ for models with $\hat{\mathcal{U}}$ , $0$ otherwise
$h_{y}$	$0.05$
$h_{x}$	N/A
Epochs	2001
$G_{\theta_{G}}$ learning rate	$5\times 10^{-3}$
$G_{\theta_{G}}$ MLP hidden layer sizes	$[128,128]$

Appendix

Table 8: Hyperparameters for semicircle experiments in

\lx@sectionsign

5.1 .

Hyperparameter	Value
$\alpha$	$0.01$ for models with $\hat{\mathcal{U}}$ , $0$ otherwise
$h_{y}$	$1.0$
$h_{x}$	$1.0$
Epochs	$2001$
$G_{\theta_{G}}$ learning rate	$5\times 10^{-3}$
$G_{\theta_{G}}$ MLP hidden layer sizes	$[128,128]$

Appendix

Table 9: Hyperparameters for our surrogate models in the cancer therapy experiment in

\lx@sectionsign

5.2.1 and

\lx@sectionsign

5.2.3 .

Hyperparameter	Value
Epochs	$10000$
Early stopping patience	$100$
Validation set	$10\%$
Batch size	$256$
Learning rate	$10^{-3}$
Hidden layer sizes	$[64,64,64,64]$

Appendix

Table 10: Hyperparameters for the direct surrogate model in the cancer therapy experiment in

\lx@sectionsign

5.2.1 .

Hyperparameter	Value
Epochs	$10000$
Early stopping patience	$100$
Validation set	$10\%$
Batch size	$1000$
Learning rate	$10^{-3}$
Hidden layer sizes	$[64,64,64,64]$

Appendix

Table 11: Hyperparameters for the CFM surrogate model in the cancer treatment experiment in

\lx@sectionsign

5.2.1 .

Hyperparameter	Value
Epochs	$2000$
Early stopping patience	$100$
Validation set	$10\%$
Batch size	$128$
Interpolant learning rate	$10^{-4}$
Vector field learning rate	$10^{-3}$

Appendix

Table 12: Hyperparameters for the MFM surrogate model in the cancer treatment experiment in

\lx@sectionsign

5.2.1 .

Hyperparameter	Value
$\alpha$	$0.001$ for models with $\hat{\mathcal{U}}$ , $0$ otherwise
$h_{y}$	$2.0$
$h_{x}$	$1.0$
Epochs	$2001$
$G_{\theta_{G}}$ learning rate	$5\times 10^{-3}$
$G_{\theta_{G}}$ MLP hidden layer sizes	$[128,128]$

Appendix

Table 13: Hyperparameters for reacher experiment in

\lx@sectionsign

5.2.2 .

Hyperparameter	Value
$\alpha$	$0.01$ for models with $\hat{\mathcal{U}}$ , $0$ otherwise
$h_{y}$	$1.0$
$h_{x}$	$1.0$
Epochs	$1001$
$G_{\theta_{G}}$ learning rate	$5\times 10^{-3}$
$G_{\theta_{G}}$ MLP hidden layer sizes	$[128,128]$

Appendix

Table 14: Hyperparameters for ETT experiment in

\lx@sectionsign

5.3 .

Hyperparameter	Value
$\alpha$	$0.01$ for models with $\hat{\mathcal{U}}$ , $0$ otherwise
$h_{y}$	$0.2$
$h_{x}$	$1.0$
Epochs	$2001$
$G_{\theta_{G}}$ learning rate	$5\times 10^{-3}$
$G_{\theta_{G}}$ MLP hidden layer sizes	$[128,128]$

Appendix

Table 15: Hyperparameters for dropout experiment in

\lx@sectionsign

5.4 .

Panel of Figure 5 — Figure 5: Surrogate model reward in Cancer (left) and Reacher (right).

Neural Surrogate Modeling

Neural surrogate modeling uses learned models to approximate costly simulations, enabling faster prediction, design, control, and inference.

See all 410 papers

Bandits in Prod: Hyperparameter Optimization at Inference Time
Sep 1, 2026Louis Abraham, Tuan-Anh Nguyen, Nicolas DevatineMulti-Armed Bandits Inference-Time Optimization
Tiime Paris, France
Quantifying Hyperparameter Transfer and the Importance of Embedding Layer Learning Rate
May 20, 2026Dayal Singh Kalra, Maissam BarkeshliLanguage Model Scaling Laws Hyperparameter Transfer
Department of Computer Science, University of Maryland, College Park · Department of Physics, University of Maryland, College Park · Joint Quantum Institute, University of Maryland, College Park +1
Hyper Input Convex Neural Networks for Shape Constrained Learning and Optimal Transport
Apr 29, 2026Shayan Hundrieser, Insung Kong, Johannes Schmidt-HieberInput Convex Neural Networks Neural Network Approximation Theory
*Department of Applied Mathematics, University of Twente, Enschede, The Netherlands

Sep 1, 2026J/K move · Enter open · S save

Bandits in Prod: Hyperparameter Optimization at Inference Time

Louis Abraham, Tuan-Anh Nguyen, Nicolas Devatine

Tiime Paris, France

Multi-Armed Bandits Inference-Time Optimization Hyperparameter Optimization Bandit Learning

Open paper AskPDF arXiv

Many production systems can assess a configuration only by using it on live requests and observing noisy feedback. Modern agentic systems are a prominent example, with inference-time choices such as model selection, retrieval depth, prompting strategy, and decoding temperature, yet often with no representative validation data. We formalize this setting as Online Hyperparameter Optimization (OHPO) and cast it as an infinitely many-armed bandit over mixed and conditional search spaces. We introduce IMABO, a general framework that combines any bandit policy for choosing among already sampled configurations with any oracle for proposing new ones. We instantiate it with IMOSS, a restart-free anytime policy whose active set grows as

t^β

, and prove an expected cumulative quantile-regret bound of

O(p_ρ^{-1/β} + T^{(1+β)/2})

, where

β\in(0,1)

controls active-set growth and

p_ρ

lower-bounds the probability that a proposed configuration falls in the top-

ρ

fraction of the search space. We combine IMOSS with three practical oracles: a Tree-structured Parzen Estimator, an incumbent-mutation oracle driven by a per-coordinate bandit, and a pretrained tabular foundation model, all three improving over the uniform random oracle baseline. IMABO outperforms all baselines in terms of regret across diverse OHPO settings, from tuning classical machine-learning models to configuring LLM-based agents. Our implementation is available at https://github.com/Tiime-Software/IMABO.

Abstract

Neural Surrogate Modeling

Explore similar work

Bandits in Prod: Hyperparameter Optimization at Inference Time

Quantifying Hyperparameter Transfer and the Importance of Embedding Layer Learning Rate

Hyper Input Convex Neural Networks for Shape Constrained Learning and Optimal Transport