Hyperparameter Trajectory Inference with Conditional Lagrangian Optimal Transport
Authors: Harry Amad
Mihaela van der Schaar
Venue: ICLR 2026 Oral (2026)
Keywords: hyperparameter
optimal transport
trajectory inference
manifold learning
interpolation
Primary area: other topics in machine learning (i.e., none of the above)

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.

Published as a conference paper at ICLR 2026
HYPERPARAMETER TRAJECTORY INFERENCE WITH CONDITIONAL LAGRANGIAN OPTIMAL TRANSPORT
Harry Amad, Mihaela der Schaar Department of Applied Mathematics and Theoretical Physics University of Cambridge, UK hmka 30 cam. ac.uk
ABSTRACT
Neural networks (NNs) often have critical behavioural trade-offs that are set at design time with hyperparameters--Sute ing 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 (HTT): learn, from observed data, how NN's conditional output distribution changes with its hyperparameters, and construct surrogate model that approximates the NN at unobserved hyperparameter settings. HTI requires extend- ing 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-induen 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.
as reward weights in reinforcement learn-
dynamics along with the associated
INTRODUCTION
Neural network (NN) behaviour is critically shaped by hyperparameters, A, which alter the parameters of the trained network, 0x thereby affecting the distribution of outputs given input x, Pe, (y|æ²) [ Often, hyperparameters govern subjective trade-offs, requiring users to fix complex preferences at design time. In deployment, however, evolving conditions can render initial hyperparameters subop- timal, necessitating retraining, which can be infeasible. This motivates an alternate approach- to learn surrogate model that can sample outputs across spectrum of hyperparameter settings. We introduce Hyperparameter Trajectory Inference (HTI) inspired by trajectory inference (TI) (Hashimoto al. 2016 Lavenant et al. 2021) to address exactly this. The goal of HTI is to learn hyperparameter induced dynamics À + Pe, (y/ar between observed distributions (pe, (3)(z) AeA.. and develop surrogate model p(y/ az, A) with which the NN conditional probability paths, for some reasonable hyperparameters E A, can be estimated as (p(y)ı,, )))xEA 2 (po, )xeA, permitting approximate inference-time adjustment of A. Below we expand two potential cases of HTT.
Reinforcement learning. NN-based RL policies (Zhu et al. 2023 Park et 2025) define a state-conditional action distribution Po, (als), with fundamental behaviours determined by certain hyperparameters. Consider, for instance, policy for cancer treatment, with a reward function balancing two objectives: reducing tumour volume, and minimising immune system damage, weighted by a scalar X. The ideal balance can vary per patient, based on factors such comorbidities (Sarfati et al.) 2016). An HTI surrogate policy p(a) X) would allow for personalised treatment strategies, by varying À at inference time (5..2.1.
'po, (wl=:) (x) (v) in the deterministic case, but also consider other distributions e.g. generative models, distributions parameterised by NN outputs.

---

Published as a conference paper at ICLR 2026
Quantile regression. Regression tasks can require measures of uncertainty. Quantile regression (Koenker & Bassett Jr. 1978) provides a way to construct prediction intervals, but typically models target individual quantiles т, or a multi-head model outputs a fixed set of quantiles (Wen et al. 2018). This can make examining arbitrary quantiles, to tailor uncertainty bounds, computationally intensive. HTI can address this, learning the dynamics desired quantile range, yiclding a surrogate that can predict all intermediate quantiles ($5.3).
ps_ (y/zz) across a
HTT is challenging, as the dynamics À + Pøs (yla:) are typically non-linear, given complex deep
learning optimisation landscapes (Ly & Gong. 2025), making simple interpolation schemes, like
conditional flow matching (CFM) (Lipman et al. 2023 Liu et al. [023] Albergo & Vanden-Eijnden
2023), unlikely to yield feasible surrogates (p(g) x. X YEA- HTI requires an approach capable of
capturing complex, non-Euclidee Euclidean dynamics from sparse ground-truth distributions. Similar problems
have been addressed in standard Ti (Tong et al. 2020 [Scarvelis & Solomon 2023; Kapusniak et al.,
2024 Pooladian et al.. 2024), however the effects of conditions on probability paths, which is essential for HT,, are currently under-explored.
We aim to enable HTT by addressing this problem of conditional TI (CTI). We propose an approach grounded in conditional Lagrangian optimal transport (CLOT) theory, allowing us to bias inferred conditional probability paths remain meaningful. Specifically, we aim learn kinetic and potential energy terms that define a Lagrangian cost function, and to encode inductive biases into these terms. This cost function determines what is deemed to be efficient movement between {pe, )EA.. and we use neural approximate solutions to the optimal transport maps and geodesics that respect this Lagrangian cost to infer conditional probability paths. We do so in a manner inspired by PPollddaan al. (2024), extending this method to handle conditions, encode more useful inductive biases, and perform on more complex and higher-dimensions components are learned, samples for target hyperparameter Auareet A and condition x can be drawn by sampling from a base distribution in (pe, (y|z) ( XeAm. approximating CLOT maps and geodesic paths, and evaluating the paths at the Ararget position. In short, our main contributions include:
geometries. Once the Lagrangian and CLOT
1. We introduce the problem of Hyperparameter Trajectory Inference to enable inference-time NN behavioural adjustment, using the framing of TI to encourage particular inductive biases for modelling hyperparameter dynamics (22.T).
2. We propose a general method for CTI to learn complex conditional dynamics from temporally sparse ground-truth samples, based on principles from CLOT (4)) We extend the procedure Pooladiian et al. (2024) in several ways, learning a data-dependent potential energy term U alongside a kinetic term KC ($4.I). elevating the method the conditional OT setting ($4.2). and
establishing a more expressive neural representation for the learned metric, Go., underpinning
K" that naturally extends higher dimensions (44.3).
We demonstrate empirically that our approach reconstructs conditional probability paths better than alternatives in multiple applications of HTI, enabling effective inference-time adaptation of a single hyperparameter in various domains (55)
PRELIMINARIES
2.1 HYPERPARAMETER TRAJECTORY INFERENCE
TI Hashimoto et al., 2016 Lavenant et al. 1021 aims to recover the continuous time-dynamics
t Pt of a population from observed samples from a set of temporally sparse distributions {pr heaa.
CTI is an extension of TI where conditioning variable x E X' affects these dynamics, with a goal of inferring the conditional population dynamies t + Pe(-Iar for arbitrary x.
Building upon the concept of CTI, we introduce a novel instantiation that we address in this work-H'TI. In HTI, the 'population' is the outputs of a NN, with distribution po, (ylz) conditioned on its input I, and we wish to learn the dynamics À - Po, (uf2) induced by single continuous hyperparameter À E A (acting as 'time") ) from a set of known distributions {pe, J Aee... to recover the conditional probability paths (po, (y)x)))xAA- These dynamics can be used to build surrogate model p(y| X) for the NN in question, allowing efficient, approximate sampling at arbitrary hyperparam- eter settings within A. Since many hyperparameters, by virtue of their effects during NN training, define families of NNs among which the optimal member is context dependent, such a surrogate model could reduce the need retrain NNs in dynamic deployment scenarios.
2

---

Published as a conference paper at ICLR 2026
2.2 CONDITIONAL OPTIMAL TRANSPORT
We deploy the framework of conditional optimal transport (COT) (Villani et al., 2008) to define optimal maps and paths between conditional distributions, which be neurally approximated. Let Vo and 111 be two complete, separable metric spaces, and X' be a general conditioning space. For X, consider probability measures [o(((z)) E P())) and (1(-1z) P(Y)) and cost function e(.--(z): Vo x 1/1 + Ryo- The primal COT formulation (Kantorovich. 1942) involves a coupling π that minimises the transport cost:
/ c(yo, yn (z) (Vo, y1 (1)2))
inf COT=(HO(-1r))) 111 (+1=)) =(vll))en(e (x)))n yox
(1)
where [-(o(-)))), mn (+rr)) is the collection of all probability measures 3' 3 with marginals po(-)æ)) on Vo and μi (-1z) on V/, Solving this primal problem is generally intractable, and it cannot
be easily neurally approximated as it requires modelling high-dimensional joint distribution. The equivalent dual formulation simplifies the problem to a constrained optimisation over two scalar potential functions [(-(2)) and g-((2r):
COTx(po(112). p1 (-1x)) = sup / ((ool=))ddo(ollz) + /,, ((s1lz)duu1(1
(2)
fag
subject to the constraint (voll=) + g(ya (x:) c(yo, 2n [x), V(uo, 3n) E (Db,)1) D1). Enforcing this
constraint with neural instantiations of o and across the entire domain is challenging (Seguy et al.
2018). As such, we follow recent literature (Makkuva et al.) 2020 Amos 2023 Pooladian al.
2024) and utilise the semi-dual formulation based on the c-transform (Villani et al. 2008 converting the problem into an unconstrained optimisation over single potential =x:
sup ((ll))E(((x((())) / g ((olz))dHo(bolz) + /
COT.(Po(-1z)).)) (+(a)) =
y. glvn (z) (3)
where (-(x) is the c-transform of g(-lz)):
(
inf (c(yo, g(vi [x)]- E/
Denoting g"(1xz) an optimal potential for 3. the COT map (-(((x)) Vo + + can be found as Te()olz) E argmin [c(vo. v((=) g" (ví lz))
vi
2.3 CONDITIONAL LAGRANGIAN OPTIMAL TRANSPORT
In the above, the cost function is where knowledge system dynamics can be embedded, to shape the COT maps and paths (Asadulaev et al. 2024). The standard Euclidean cost c(Uo- 9/) = ||sto V/ 1| for example, deems straight paths as the most efficient. To induce more complex paths, given our assumed complex syperparametee-dynaaiiii motivates us to use principles from Lagrangian dynamics (Goldstein et al. 1980), bringing us the CLOT setting. Given a smooth, time-dependent curve qe for l. E [0, 1J. with time derivative dt and Lagrangian C(ge. gelz), the action q. S(q(z), can be determined as
we require a cost function that is path-dependent, which
S(glx) = S( L(qe, de delz))t.
(6)
The resulting Lagrangian cost function c can then be defined using the least-action, or geodesic, curve between Yo and
c(yo- V1 9:90 0o yn
inf S(q.z).
We denote geodesics as q². While flexible in form, common Lagrangian instantiation is
C(g,, de (z) = K(ge, (x) - ((]. = ot G(qll)4 - ll(g(())
(8)
2
where IC and 24 are kinetic and potential energy terms, respectively, with metric G defining the geometry of the underlying manifold (e.g. for Euclidean manifolds, G -= /). We consider learning conditional Lagrangians of the above form by setting a neural representation of G and estimating U using a kernel density estimate, and learning neural estimates of the transport maps and geodesics for the consequent CLOT problem. We design and G to incorporate biases for dense traversal and least-action into the inferred conditional probability paths.
3

---

Published as a conference paper at ICLR 2026
3 RELATED WORKS
Trajectory inference. TI (Hashimoto et al. 2016 Lavenant al. 2021 is prominent in domains such as single-cell genomics, where destructive measurements precluc tracking individual cells over time (Macosko et al., 2015 Schiebinger et al. 2019). Successful TI relies on leveraging inductive biases to generalise beyond the sparse observed times. One typical bias is based on least-action principles- assuming that populations move between observed marginals in the most efficient way possiblc-naturallly giving rise to OT approaches (Yang & Uhler 2019; Scciieiinger et al.) [019 Tong et al. 2020, Scarvelis & Solomon 2023 Pooladian et al. 2024.. Another potential bias invokes the manifold hypothesis (Bengio et al. 2013), which posits that data resides on a low-dimensional
manifold, concentrated around the observed data (Arvanitidis et al. 2022 Chadebec & Allassonnière 2022). to encourage inferred paths to traverse dense regions of the data space (Kapusniak et al. 2024).
Neural optimal transport. NNs have been used for OT, especially in high-dimensions where classical OT algorithms infeasible (Makkuva et al.) 2020; Korotin et al. (022I) The semi-dual OT formulation with neural paramctrisations of the Kantorovich potentials and transport maps is
standard (Makkuva et al. 2020; Amos 2023] [Pooladian et al. 2024). Neural COT has also been explored Waag et al. 2024 2025), although with fixed cost functions, and we novelly extend this to incorporate learned conditional Lagrangian costs. Our work is particularly related to [Scarvelis) & Solomon 2023) and Pooladian et al. (2024), which jointly learn OT cost functions and resulting transport maps from observed time marginals. We consider using more expressive forms for the cost function, involving Lagrangians with kinetic and potential energy terms, and we operate in the conditional OT setting.
Conditional generative modeling via density transport. Some generative models, such as condi- tional diffusion (Ho & Salimans, 2022) and CFM models (Zheng et al., 2023), operate by transporting mass from a source to a target distribution, according to some condition. They can therefore be applied to conditional TI. However, generative models focus on accurately learning the target data distribution, and they are generally unconcerned with the intermediate distributions formed along the transport paths. While some recent works utilise OT principles to achieve more efficient learning and sampling for CFM models (Tong et al., 2024 Pooladian et al. 2023). their primary objective remains high-fidelin sample generation from the target distribution.
Bayesian optimization. Bayesian hyperparameter optimization (Snoek et al. 2015) builds a surrogate model of a NN's objective function across hyperparameters. HTT extends this by learning surrogate for the NN's conditional output distribution rather than for a scalar objective function. HTI could allow for more flexible Bayesian hyperparameter optimisation with arbitrary, post-hoc objective functions calculated with surrogate samples (Appendix Ap.
2012 Shahriari et al.,
4 NEURAL CONDITIONAL LAGRANGIAN OPTIMAL TRANSPORT
We now present our method for general CTI, which involves neural approach to CLOT. From ob- served temporal marginals, we seek to learn both the underlying conditional Lagrangian L(g (|x>) K(q, < Ld(q)a æ) that governs dynamics, along with the consequent CLOT maps T. and geodesics q². such that conditional trajectories can be inferred. We novelly encode both the inductive biases discussed in Section 3 lcast-action and dense rraveral--ite L to aid generalisation of inferred trajectories beyond the observed temporal regions.
4.1 POTENTIAL ENERGY TERM
Firstly, we set the conditional potential energy, (g]z), through which we encode a bias for dense traversal. By designing (q(()) to be large in dense regions of the data space, and small elsewhere, the Lagrangian cost function c, as in will lead to geodesics that favour dense regions.
Let Dots {v,, Fi, tx)) ),, be the set of observed samples, where y; y are the Dyddimensional ambient space observations, x4 E X' are their D -.-dimensional conditions, and t2 E {tu, ti.... tr the T 'times" of observation. We define the potential at q for given condition æ E X' as:
ü(q a log(ū(g(>)) x:) + e),
4

---

Published as a conference paper at ICLR 2026
where 0 is set by the user to control the strength of the density bias, 0 is for numerical stability, and p(q(z) is estimated with Nadaraya- Watson estimator (Nadaraya 1964 Watson 1964):
((gaa) Σ/ Ki, (q V6 K (x (x- xz))
(10)
Σラ=1 Kr_ (x xz)
where Ki, and Ka, are Gaussian kernel functions with bandwidths By and hag respectively:
Kny (u) (2Th exp (- [||u112)
||:||
D2/2
Kn_ (v) (2mh2)) Dz/2 exp
(11)
2h
2h
We that (9) will be high when p(glæ)) high, and low when p(g]ar) is low, thus encoding our desired bias for geodesics traverse dense regions of the data space. [gqga)) fixed throughout the subsequent learning phase for the kinetic energy term K. and the CLOT maps and geodesic paths.
4.2 JOINT LEARNING OF KINETIC ENERGY TERM AND CLOT PATHS
To learn the remaining kinetic term KG d|z) = 3r" q. and solve the consequent CLOT problem, adopt neural approach similar to Pooladian al. 2024), adapting it to conditional setting. We operate under the assumption that the observed data display dynamics that are efficient in the underlying data manifold, to embed the desired least-action bias into method. We consider neural instantiations of the metric Gag within IC, and the 7 Kantorovich potentials go,. defining the CLOT problems between temporally adjacent observed distributions, with parameters 0g and 0g.k respectively These networks are learnt with a min-max procedure, alternating between optimising Ge.. with fixed .e,.k. to minimise the estimated CLOT cost between observed marginals (encoding the desired least-action principles), and optimising each g8,.k, with fixed Gag. maximise (3) (to accurately estimate the CLOT cost under the current metric). The overall objective is
min AG
Σ Fez
max (- 19s,,. (3n] + E ( 100 (yk+1 ( Pa.k
(12)
where pa (-|2z) is the conditional distribution of the data time 44. We denote the inner maximisation objective for each interval as (k) (0g.k), and the outer minimisation objective as Emerii (fc).
Calculating 9², as in (4) requires solving an optimisation problem, with further embedded optimi- sation problem to calculate the transport cost. These nested optimisations can make training computa- tionally infeasible. As such, adopt the amortisation procedure from Poolaadian et al. (2024), simul- taneously training and using neural approximators to output CLOT maps Ter, (3k R Te (VA (z) and the parameters approximation. At given training iteration, the current learned map Te warm-starts the minimisa- tion ( this estimate is refined with limited number of L-BFGS (Liu & Nocedal 1989) steps to yield Te, Ty, (x) which is used to calculcate g in (12), and regression target for Ter.
spline-based based geodesic estimation, Qg q"' allowing efficient e-transform
mmpp (0T, (s) )== E [(Tor, (yn(lr)) (x)))())))))]
(13)
To efficiently calculate the cost function required for these L-BFGS steps we approximate geodesic paths qº with cubic spline qp- with parameters &p output by NN Sog trained to minimise
Spahh (0s) E [S(q,\=)] & = Sa (yk: Tor, (3x ()).
(14)
To condition each network x, equip them with FiLM layers (Perez et al. 2018) that modulate the first-layer activations based on æ. The overall training procedure (Algorithm) 1) alternates between
updating each Ky fT and 0 to maximise the inner part of (12). minimise (13). and minimise 11.
respectively, and updating 0g to minimise the outer sum in (2)
4.3 METRIC PARAMETRISATION
Within the above procedure, the parametrisation of the neural metric Gog particularly important, as this must be symmetric, positive-defiitt, D,yddimensional matrix to be valid metric. Critically,
9o,,, n denotes the th Kantorovich potential, for the CL.OT between the distributions f3 and f++1
5

---

Published as a conference paper at ICLR 2026
Algorithm 1 Neural CLOT Training
Require:
Observed data Dobs, ambient and conditional bandwidths by, ft,, potential weight no. outer training iterations Nouter. no. inner training iterations Ninners learning rates fgg (r, Vls 1J
1: ((gzz) + l)g( where ò(g/z=) Σ; Khn, (9- 1--1i) (n))aa(==E))
Σi=ı Kng (ar (x-2z)
2: Initialise lc. g0 0T tR. Os
3: Define S(q = ( (247 Goo (q:1=) ((ql(2))d
4: for i = 1... Nowter do
5: for j = 1..... mmerr do
6: for k 0... T - do
7:
D& t {,,, xt E Dobs -t=
8:
for (yk: x.) E Dr. do
9:
Uk -- Tor, (yx:3))
10:
y/ + L-BFGS(y) S(1el()) [z) 90,, (V(lz)), where Φi = S8g (yk the
11:
9o (yx|z) + S(go- (x) - 10,. (2 (2)), where o² = Sos (Jles 2k x0)
12:
end for
13:
gg.k f- 0g. y 7g V C (k)
dual (0g.b) 0r + Br. 7/2 V Vmmap(01 (0
14:
end for
15:
0s 8g ns V path (8s)
16:
end for
17:
0G 0g N| V mmerrii (fG))
18:
end for
return (Ter, J&, Sag
there exist degenerate minima to (12) by setting Goe + 0. where movement all directions results in near-zero cost. We set our parametrisation to ensure Ge avoids this and maintains sufficient volume. In Pooladian et al. (2024), where only two-dimensional data spaces are considered, they set Geo fixed diagonal matrix with neural rotation matrix
Tcos( Reg (x)) sin Ro (x))
sin(Rec (x)) cos( Ro (I))
1 0
0
cos Roc (I)) sin Rocs (x)) cos(Rea (I))
Rin (x))
Goo
(15)
0.1
where Rea (x) is the output of the NN. This is only applicable to two-dimensional spaces, and avoids degenerate solutions by fixing the local anisotropy of Gng. We design a parametrisation that extends to higher dimensions, and is more expressive, while still avoiding degenerate solutions without requiring regularisation as in Scarvelis Solomon (2023). Specifically, we set G&o using its cigendecomposition Gs Rg Egg where NN parametrises both Dy-dimensional diagonal matrix Fag- and rotation matrices Ro.- To avoid degeneracy, we enforce the entries of Eog- and therefore the eigenvalues of Gog- to be positive, and sum to non-zero 'eigenvalue budget', ensuring non-trivial volume of Goo while permitting expressive levels of anisotropy. To define
the dyddimeesional rotation matrix Rag. we multiply (D, Givens rotation matrices (Givens 1958), with angles parametrised by the NN. This can improve performance over the fixed approach of Pooladian al. (2024) in two-dimensions (55.)) while also extending to higher dimensions ($5.3)
4.4 SAMPLING ALONG THE INFERRED TRAJECTORY
To generate samples from the inferred conditional distribution P(y|2, t²), use the neural approxi- mators for the CLOT maps and geodesics, avoiding the need for any optimisation at inference time. First, samples are drawn from the ground-truie distribution with the largest observed base time with t < t", Yk ~ Pt. (-|zz). ) The learned map Tor, (yk (z:) then predicts the transported point yk+1 at the end of the interval [ta. fx+1]. which contains t". Subsequently, the parameters for the approximate geodesic path 9, connecting ys to yk++1 can be estimated as ip -e,, (yk. Ms+1. and q, can be evaluated at the appropriate time. By normalising t" to s* = (t" -tx) (tx/ tx), the inferred sample obtained ÿt- (8*).

---

Published as a conference paper at ICLR 2026
5 EXPERIMENTS
We now empirically demonstrate the efficacy of method for CTI and some specific HTI scenarios. All results are averaged over 20 runs, and reported with standard errors. We provide detailed experimental set-ups in Appendix C
5.1 ILLUSTRATIVE EXAMPLE OF CTI
To illustrate our method's inductive biases, we devise temporal process with conditions x E {1, 213 4}, where each defines distribution (e(z))) that evolves from the origin over é E [0, as noised von Mises, with centre moving along one four semicircular paths. Samples from the true process shown in Figure [ where each condition has a distinct colour, and lighter samples are from larger To conduct CTI, using observations from E {0, 0.5, 1.0]. models must: (1) learn condition-dependent dynamics despite overlapping initial distribu- tions; (2) capture the -on-Euclidean geometry of semicircular paths; and (3) generalise across t [0, 1] from sparse temporal samples.
(a) Kr
(b) Kr
(c) Ko
(d) Kg - ú
Figure Dots represent true samples from the temporal process across t [0. 1]. lines represent
model estimated trajectories from t = 0 to Each condition has distinct colour.
We compare four ablations of our method, with varying complexity of the learned conditional Lagrangian: (1) Ky; Using an identity metric G = and setting ú = O, resulting in Euclidean geometry with no density bias; (2) Kg:: Learning the metric Gg., via our method in $4.2 and setting
U = 0, to incorporate only the inductive bias of least-action; (3) Kr - út Using an identity metric
G = I and learning Li as in $44. to incorporate only the inductive bias of dense traversal, and; (4)
Kg L7: Our full approach, learning both the metric Gae and the potential term U.
Figure 0 shows the inferred paths of samples from t t= 0 to 1. Our full method (Figure lcl most faithfully reconstructs the true temporal process, as the paths correctly diverge according to their con- dition and closely follow the intended semicircular geometry. We can see the individual effects of both inductive biases, individually learning úí (Fig- ure Tb) results in straight paths that favour denser regions, avoiding the circle centres, while learning
Table NLL and CD / (0.25, 0.75).
Method
NLL (L)
CD (L)
K₁
105.713 (2.42)
0.323 (0.003)
Ka
23.008 (4.62)
0.158 (0.009)
KG, --
-0.532 (0.057)
0.016 (0.001)
Ko -u -0.662 (0.046) 0.016 (0.001)
Gog only (Figure Ic] better captures the underlying curvature of the semicircular geometry. In
Table I we evaluate plylæ, J at withheld e {0.25, 0.75), reporting negative likeliiod (NLL) and distance from the target circle perimeter (CD). We can quantitatively see that both inductive biases improve the feasibility of the inferred marginals.
5.2 HTI FOR REWARD- WWARR--WIGGHTIGG REINFORCEMENT LEARNING
We now transition to specific applications of HTI, first addressing compelling challenge in RL, to create surrogate policies that allow for dynamic reward weighting.
5.2.1 CANCER THERAPY
We investigate HTI for personalised cancer therapy, mirroring the first use case presented in | We employ an environment from STR-Bench (Luo et al. 2024), which we call Cancer, that simulates tumour progression under chemotherapy and radiotherapy. Natural Killer (NK) cells are pivotal immune system components, and they be depleted as side effect of cytotoxic treatments like chemotherapy and radiotherapy (Shaver et al., 2021] Toffoli et al. (202I), increasing susceptibility to infections and compromising treatment effiaay. This side effect varies substantially with age,
7

---

Published as a conference paper at ICLR 2026
Table 2: Average surrogate Cancer reward across Ank (1, 4, 6, 9,9.
Thue
0.06
HTI Surrogate
Pwn
Method
Reward (t)
0.04
Direct NLOT
38.35 (10.65)
9.26 (10.55)
K.s
30.63 (8.50)
0.02
CFM
36.03 (6.46)
0
2
6
MFM
41.05 (4.16)
ANc
Kr
48.72 (7.22)
K1 -u U
83.62 (5.37)
Figure PNKK vs. Ank for ground truth policies and our surrogate policy.
Ko -u
102.49 (5.46)
comorbidities, and baseline immune status (Diakos et al. 2014) and, consequently, optimal cancer therapy necessitates patient-specific balance between tumour reduction and NK cell preservation.
The Cancer reward function incorporates both tumour volume and NK cell preservation, with hyperparameter An weighting an NK cell penalty term, PNK. Training a Proximal Policy Optimiza- tion (PPO) (Schulman al. 2017) agent to convergence in this environment takes approximately 3.5 hours, so training per-patient policies with tailored Ank is computationally prohibitive. This therefore presents prime application for HTI, to enable inference-time policy adaptation.
To learn the Aax-induced dynamics of the policy distribution across Auk E [0, 10), we train ground- truth policies with PPO Ank E {0, 5, 10} and sample 1000 state-action pairs from each converged policy, across a shared set of states, to act as the HTI training set. We assess the four approaches from [5.1] alongside some baselines. We compare to: (1) direct surrogate, where the target Aak current state, and action from Ank = 0 are inputs to an MLP that is trained to output actions at given Ank via supervised regression; (2) CFM surrogate (Lipman al.. 2023), which learns vector field between the distributions at Ank E (0. 5, 10) and generates samples by integrating actions to the desired Ank point; (3) a metric flow matching (MFM) surrogate (Kapusniak et al.) 2024) that is similar to CFM, but biases the vector field to point towards dense regions of the data space; and (4) the NLOT method of Pooladian et al. (024). We add FiLM conditioning to these to make them appropriate for HTI.
In Table 2 we report the average reward for each surrogate at held-out Auk E 2, 4,6, 7, 7, 9). Our full method (Ko 27) infers the most realistic trajectory between settings, yielding surrogate policy with the best average reward. We also examine NK cell preservation in Figure plotting the average per-episode PNK penalty for our surrogate and ground-truth policies. Our method closely mirrors the behaviour of the ground-truth policies, correctly favouring treatment that preserves NK cells as Ank increases. Critically, training our surrogate model takes approximately 15 minutes, after which rapid inference-time adaptation is possible. This contrasts with the 3.5 hours required to train each new PPO policy, highlighting the substantial computational advantage conferred by HTI.
5.2.2 REACHER
Table 3: Average surro-
To further demonstrate HTI for reward weighting, we evaluate in the Reacher environment from OpenAl Gym (Brockman et al. 2016), standard continuous control benchmark. In this setting, an agent controls two-joint arm with the goal of reaching random target position. The reward designed to minimise distance to the target, while penalising the magnitude of the joint torques, discour- aging high-force movements, and it is weighted by hyperparameter X.
gate Reacher rewards across A E 32,3,4).
Method
Reward (t)
Direct MFM
6.711 (0.070)
-6.561 (0.053)
Kr, ú
-6.397 (0.031)
KG,
-6.307 (0.041)
Similar to the cancer therapy experiment, we train PPO agents at d E {1, 5), and collect 1000 state-action pairs from each agent to form the HTI training dataset. In Table Ba we evaluate the same suite of surrogate models previously, assessing inferred policy
CFM
-6.251 (0.028)
NLOT
-6.173 (0.038)
Ko
-6.158 (0.033)
Ko - ú
-6.093 (0.036)
8

---

Published as a conference paper at ICLR 2026
behaviour unseen X (2, 4). Our full method (Coo - 2t) again yields the most performant surrogate, achieving the highest average reward.
5.2.3 NON-LINEAR REWARD SCALARIZATION
Table 4: Average surrogate
The previous reward scalarizations involve linear combinations of a main objective (tumour volume/distance to target) and a penalty term (NK penalty/torque penalty). Such scalarization is known lead to well-behaved trade-offs when tuning reward weights (Rä- dulescu et al., 2020). For a challenging RL setting, with less well-behaved hyperparameter dynamics, we modify Cancer to have non-linear reward scal scalarization, with hinge penalty. In this Cancer_nd setup, the weighted NK penalty is only applied the change in cell count crosses threshold (see definition Ap- pendix C.2.2 We employ the same training and evaluation protocol as in 5.2.1 with results Table We see that our method again achieves the highest average reward across held-out settings, remain- ing robust when the hyperparameter governs non-linear objectives.
Cancer n1 reward across
Ank E (1, 8,,
Method
Reward (t)
K1 Direct
42.84 (5.86)
49.50 (17.90)
NLOT
52.63 (6.52)
Kaa
54.44 (7.66)
CFM
69. 69.70 (7.73)
MFM
78.44 (4.47)
KG ú
86.21 (6.32)
Kg -u ú
91.94 (11.46)
5.3 HTI FOR QUANTILE REGRESSION
We now demonstrate HTT's application in a higher-dimensional setting of quantile regression for time-series forecasting, mirroring the second use case presented in [ Time-series forecasting is a task where providing a full picture of uncertainty, such as through quantile regression, is crucial, but the need train forecasting models target distinct quantiles can hinder this. We investigate whether HTI can address this by inferring intermediate quantiles from the outputs of models trained at the extremes the quantile range. Using the ETTm2 forecasting dataset (Zhou al. 2021), we train two MLPs to forecast 3-step horizon from 12-step history the quantiles using a standard pinball loss. then generate dataset of 1200 forecasts from these two models, across shared inputs, to act as the HTI training set. In Table 5 we evaluate the mean squared error (MSE) for surrogate forecasts at held-out quantiles T E (0.1, 0.25, 0.5, 0.75, 0.9} compared to the true NN outputs unseen input data. Our full method once again outperforms all baselines. Figure provides qualitative validation of this, visualising the central 80% prediction intervals (between the 0.1 and 0.9 quantiles) from different surrogates on random selection of samples, alongside the 80% intervals from the ground-truth NNs. Our method most closely matches the width and shape of the true interval.
0.01 and 7 0.99,
Hultory Thue
Direct
CFN
MFH
Ours
Table 5: MSE of surro- gate ETTm2 forecasts compared to NNs trained across quantiles (0.1, 0.25, 0.5, 0.75, 0.9).
Method
MSE (L)
Direct
1.845 (0.065)
0.5
CFM
1.402 (0.008)
63
MFM
1.387 (0.022)
Krr
0.765 (0.070)
K2 ú
0.651 (0.076)
Central 80% prediction intervals from HTT surrogates compared with the true intervals randomly selected ETTm2 samples, for direct (top), CFM (second row), MFM (third row), and our (bottom) approach.
Figure
K
0.620 (0.057)
Ko
ú
0.608 (0.034)

---

Published as a conference paper at ICLR 2026
5.4 HTI FOR DROPOUT IN GENERATIVE MODELLING
Table 6: two moons distributions from ground-truth
WD of surrogate
To demonstrate a more general use case for HTI, we consider the dropout hyperparameter, which is often used for regular- isation in NN training. We consider a simple setting, train- ing a diffusion model (Ho et al. 2020) on the sklearn two moons dataset at dropout settings E {0, 0.5, 0.99), and we draw 1000 samples from each model to act as the HTI training set. In Table 6 we evaluate surrogate models at held-out set- tings p E {0.I, 0.2, 0.3, 0.4, 0.6, 0.7,0.8, 0.9}. measuring the Wasserstein distance (WD) between surrogate and ground-truth distributions. Our methods incorporating the density bias ži perform well here, with K,2 ú achieving the lowest WD. We see that HTI can be used to interpolate across a general hyper- parameter such as dropout with minimal error in this setting.
diffusion mod- els, trained at dropout p 0.2,, 0.3, 0.4, 0.6, 0.7, 0.8. 0.9).
Method
WD (L)
K₁
0.570 (0.020)
MFM
0,480 (0.020)
CFM
0.464 (0.021)
Direct
0.450 (0.013)
Kg
0.404 (0.018)
NLOT
0.291 (0.022)
Kn ú
0.079 (0.003)
K1 ú
0.060 (0.001)
5.5 METRIC LEARNING ABLATION
In Table 7 we compare our neural metric Go... with learned rotation Reg and eigenvalues Egg, against the parametrisation from Pooladian et al.] (2024), which uses fixed eigenvalues E. We evaluate both within our most expressive Lagrangian setting across the previous experiments with two-dimensional ambient spaces. Our parametrisation leads better performance in most tasks, yielding lower NLL in the semicircle task and higher rewards in the Cancer and Reacher environments. On the other hand, it achieves a slightly worse WD in the generative modelling droput experiment. These results suggest that it is possible to learn the eigenvalues of Go and this can enable a more accurate recovery of the underlying conditional dynamics. Furthermore, unlike the parametrisation of Pooladian et al. (2024), ours readily extends to higher-dimensional settings, as demonstrated in 5.3
Table 7: Gog ablations in 2D experiments.
Semicircle
Cancer
Reacher
Dropout
Goa
NLL (L)
CD (L)
Reward (1)
Reward (1)
WD (L)
Ro ERT
0.602 (0.033)
0.016 (0.001)
98.72 (6.32)
-6.122 (0.080)
0.076 (0.003)
Reg Eng Rga
0.662 (0.016)
0.016 (0.001)
102.49 (5.16)
-6.093 (0.036)
0.079 (0.003)
DISCUSSION
In this work, we investigate CTI, proposing novel methodology grounded in the principles of CLOT. Our approach extends existing TI techniques by explicitly incorporating conditional information, and novelly combining dense traversal (via úi) and least-action (via KGa) inductive biases. Our empirical results show we can effectively reconstruct non-Fuclidean conditional probability paths from sparsely observed marginal distributions (..T). Our ablation study validates our neural metric parametrisation, highlighting its ability to capture intricate data geometries (55.5) while extending to higher dimensions (..3). We also investigate performance at different sparsity levels in Appendix E Furthermore, we propose HTI as novel and impactful instantiation of CTI, addressing the challenge of adapting NN behaviour without expensive retraining. We showcased the practical utility of HTI for interpolating between reward weights in RL (5..2), quantile targets in time-series forecasting (3.3). and dropout settings in generative modelling ($5.4). For reference on the efficiency HTT can confer, the ground-truth result in Figure 2 required training PPO policies, taking approximately 38 GPU hours, while the surrogate result requires training three PPO policies and an HTI surrogate, taking only 11 GPU hours. Further potential applications of HTI are discussed in Appendix A
Nevertheless, HTI will be challenging when the underlying dynamics are chaotic, making inference from sparse samples inherently difficult. Given the relatively simple settings we demonstrate here, further investigation across wider range of hyperparameter landscapes is warranted. Also, our method for HTI is only applicable for varying a single, continuous hyperparameter. Future work should explore extensions to handle multiple hyperparameters, which discuss in Appendix [)
10

---

Published as a conference paper at ICLR 2026
REPRODUCIBILITY STATEMENT
We are committed to ensuring our work is reproducible. As such, we give a brief introduction to the mathematical concepts method is based on in [ clearly describe method in $ and provide concrete training and sampling algorithms in @ respectively. For the results in [ we give detailed experimental set-ups in Appendix C and environments used, model hyperparameters and training procedures, and providing references and links to key libraries. Furthermore, we release our code base here: https:/ hett://ggtuu... com/ tra jectory-ttttioo ineennnnceeeeee
This includes detailing the datasets
ACKNOWLEDGEMENTS
Harry Amad's studentship is funded by Canon Medical Systems Corporation. This work was supported by Azure sponsorship credits granted by Microsoft's Al for Good Research Lab.
REFERENCES
Michael Samuel Albergo and Eric Vanden-Eijnden. Building normalizing flows with stochastic interpolants. In The Eleventh International Conference on Learning Representations, 2023.
Brandon Amos. On amortizing convex conjugates for optimal transport. In The Eleventh International Conference on Learning Representations, 2023.
Georgios Arvanitidis, Miguel González-Duque Alison Pouplin, Dimitris Kalatzis, and Søren Hauberg. Pulling back information geometry. In 25th International Conference Artificial Intelligence and Statistics, 2022.
Arip Asadulaev, Alexander Korotin, Vage Egiazarian, Petr Mokrov, and Evgeny Burnaev. Neural optimal transport with general cost functionals. In The Twelfth International Conference Learning Representations, 2024.
Yoshua Bengio, Aaron Courville, and Pascal Vincent. Representation learning: A review and new perspectives. IEEE transactions pattern analysis and machine intelligence, 35(88:1778-1888 1828, 2013.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. Openai gym. arXiv preprint 1rXiv: 2016.
Christopher P Burgess, Irina Higgins, Arka Pal, Loic Matthey, Nick Watters, Guillaume Des- jardins, and Alexander Lerchner. Understanding disentangling in beta-vae. arXiv preprint arXiv:1804.03599, 2018
Clément Chadebec and Stéphanie Allassonnière. geometric perspective on variational autoencoders. Advances in Neural Informaionn Processing Systems, 35:19618- 2022.
Connie Diakos, Kellie A Charles, Donald C McMillan, and Stephen J Clarke. Cancer-related inflammation and treatment effectiveness. The Lancet Oncology, 15(11 10493-e503, 2014.
Ghaffari, B Bahmaie, and M Nazari. A mixed radiotherapy and chemotherapy model for treatment of cancer with metastasis. Mathematical methods in the applied sciences, 2016.
Wallace Givens. Computation of plain unitary rotations transforming general matrix to triangular form. Journal of the Society for Industrial and Applied Mathematics, 6(1):26-50, 1958.
H Goldstein, C Poole, and Safko. Classical mechanies, 1980.
Ian J Goodfellow, Jonathon Shlens, and Christian Szegedy. Explaining and harnessing adversarial examples. stat, 1050:20, 2015.
Tatsumori Hashimoto, David Gifford, and Tommi Jaakkola. Learning population-lvvel diffusions with generative mns. In International Conference on Machine Learning, pp. 2417-2426. PMLR, 2016.
11

---

Published as a conference paper at ICLR 2026
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta-vae: Learning basic visual concepts with a constrained variational framework. In International conference learning representations, 2017.
Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12500 2022.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-688 2020.
Leonid V Kantorovich. On the translocation of masses. In Dokl. Akad. Nauk. USSR (NS), volume 37, pp. 199-201, 1942.
Kacper Kapusniak, Peter Potaptchik, Teodora Reu, Leo Zhang. Alexander Tong, Michael Bronstein, Joey Bose, and Francesco Di Giovanni. Metric flow matching for smooth interpolations on the data manifold. Advances in Neural Information Processing Systems, 37: 37:13501 1-135042, 2024.
Diederik Kingma and Max Welling, Auto-encoding variational bayes. stat, 1050:1, 2014.
Roger Koenker and Gilbert Bassett Jr. Regression quantiles. Econometrica: journal of the Economet- ric Society, pp. 33 33-50, 1978.
Alexander Korotin, Lingxiao Li, Aude Genevay, Justin M Solomon, Alexander Filippov, and Evgeny Burnaev. Do neural optimal transport solvers work? continuous wasserstein-2 benchmark. Advances in neural information processing systems, 34:14593-146055 2021.
Hugo Lavenant, Stephen Zhang, Young-Heee Heon Kim, and Geoffrey Schiebinger. Towards a mathemati- cal theory of trajectory inference. stat, 1050:18, 2021.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matthew Le. Flow matching for generative modeling. In The Eleventh International Conference Learning Representations, 2023.
Yaron Lipman, Marton Havasi, Peter Holderrieth, Neta Shaul, Matt Le, Brian Karrer, Ricky TQ Chen, David Lopez-Paz, Heli Ben- Hamu, and Itai Gat. Flow matching guide and code. arXiv preprint TrXiv:2412.06264, 2024.
Dong C Liu and Jorge Nocedal. On the limited memory bfgs method for large scale optimization. Mathematical programming, 45(1): (5(1):503528. 1989.
Xingehao Liu, Chengyue Gong, al. Flow straight and fast: Learning to generate and transfer data with rectified flow. In The Eleventh International Conference on Learning Representations, 2023.
Zhiyao Luo, Mingcheng Zhu, Fenglin Liu, Jiali Li, Yangchen Pan, Jiandong Zhou, and Tingting Zhu. Dtr-bench: an in silico environment and benchmark platform for reinforcement learning based dynamic treatment regime. arXiv preprint arXiv:2405. 18610, 2024.
Andrew Ly and Pulin Gong. Optimization on multifractal loss landscapes explains diverse range of geometrical and dynamical properties of deep learning. Nature Communications, 16(1):3252, 2025.
Evan Z Macosko, Anindita Basu, Rahul Satija, James Nemesh, Karthik Shekhar, Melissa Goldman, Itay Tirosh, Allison R Bialas, Nolan Kamitaki, Emily Martersteck, al. Highly parallel genome- wide expression profiling of individual cells using nanoliter droplets. Cell, 1222 2015.
Aleksander Madry, Aleksandar Makelov, Ludwig Schmidt, Dimitris Tsipras, and Adrian Vladu. Towards deep learning models resistant to adversarial attacks. In International Conference on Learning Representations, 2018
Ashok Makkuva, Amirhossein Taghvaci, Sewoong Oh, and Jason Lee. Optimal transport mapping via input convex neural networks. In International Conference on Machine Learning, pp. 6672-6681. PMLR, 2020.
12

---

Published as a conference paper at ICLR 2026
Elizbar A Nadaraya. On estimating regression. Theory of Probability & Its Applications, 9(1): 141-142, 1964.
Seohong Park, Qiyang Li, and Sergey Levine. Flow q-learning. In International Conference on Machine Learning (ICML), 2025.
Ethan Perez, Florian Strub, Harm De Vries, Vincent Dumoulin, and Aaron Courville. Film: Visual reasoning with general conditioning layer. In Proceedings of the AAAI conference artificial intelligence, volume 32, 2018.
Aram-Alexandre Pooladian, Heli Ben-Hamu, Carles Domingo-Enrich, Brandon Amos, Yaron Lipman, and Ricky T Chen. Multisample flow matching: Straightening flows with minibatch couplings. ICML 2023, 2023.
Aram-Alexandre Pooladian, Carles Domingo-Enrich, Ricky Tian Qi Chen, and Brandon Amos. Neural optimal transport with lagrangian costs. In Proceedings of the Fortieth Conference on Uncertainty in Artificial Intelligence, pp. 2989-3003, 2024.
Roxana Rädulescu, Patrick Mannion, Diederik M Roijers, and Ann Nowé. Multi-objective multi- agent decision making: utility-basse analysis and survey. Autonomous Agents and Multi-Agent Systems, 3401):10, 2020.
Antonin Raffin, Ashley Hill, Adam Gleave, Anssi Kanervisto, Maximilian Ernestus, and Noah Dormann. Stable-baselines: Reliable reinforcement learning implementations. Journal of Machine Learning Research, 22(268): 1-8, 2021. URL / 20 1364. html
Diana Sarfati, Bogda Koczwara, and Christopher Jackson. The impact of comorbidity on cancer and its treatment. CA: cancer journal for clinicians, 66(4):337-355 2016.
Christopher Scarvelis and Justin Solomon. Riemannian metric learning via optimal transport. In International Conference Learning Representations. OpenReview, 2023.
Geoffrey Schiebinger, Jian Shu, Marcin Tabaka, Brian Cleary, Vidya Subramanian, Aryeh Solomon, Joshua Gould, Siyan Liu, Stacie Lin, Peter Beruhe, et al. Optimal-transpon analysis of single-cell gene expression identifies developmental trajectories in reprogramming. Cell, 176(4):928-99)) 9433, 2019.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017.
Vivien Seguy, Bharath Bhushan Damodaran, Remi Flamary, Nicolas Courty, Antoine Rolet, and Math- ieu Blondel. Large-scale optimal transport and mapping estimation. In ICLR 1018-nnternatonall Conference Learning Representations, pp. 1-15, 2018.
Bobak Shahriari. Kevin Swersky. Ziyu Wang, Ryan Adams, and Nando De Freitas. Taking the human out of the loop: A review of bayesian optimization. Proceedings of the IEEE, 104(1): 148-175, 2015.
Kari Shaver, Tayler Croom-Perez, and Alieja J Copik. Natural killer cells: the linchpin for successful cancer immunotherapy. Frontiers in immunology, 12:679117, 2021.
Jasper Snoek, Hugo Larochelle, and Ryan P Adams. Practical bayesian optimization of machine learning algorithms. Advances in neural information processing systems, 25, 2012.
Elisa C Toffoli, Abdolkarim Sheikhi, Yannick D Höppner, Pita de Kok, Mahsa Bazdanpanah-Samart Jan Spanholtz, Henk MW Verheul, Hans J van der Vliet, and Tanja D de Gruijl. Natural killer cells and anti-cancer therapies: reciprocal effects on immune function and therapeutic response. Cancers, 13(4):711, 2021.
Alexander Tong, Jessie Huang, Guy Wolf, David Van Dijk, and Smita Krishnaswamy. Trajectorynet: A dynamic optimal transport network for modeling cellular dynamics. In International conference machine learning, pp. 9526-9536. PMLR, 2020.
13

---

Published as a conference paper at ICLR 2026
Alexander Tong, Kilian Fatras, Nikolay Malkin, Guillaume Huguet, Yanlei Zhang, Jarrid Rector- Brooks, Guy Wolf, and Yoshua Bengio. Improving and generalizing blowdbased generative models with minibatch optimal transport. Transactions Machine Learning Research, pp. 1-34, 2024.
Cédric Villani al. Optimal transport: old and new, volume 338, Springer, 2008
Richard von Mises. Über die "ganzzahligkeit" der atomgewichte und verwandete fragen. Physikalis- che Zeitschrift, 19:490, 1918.
Jun Wang, Bohan Lei. Liya Ding, Xiaoyin Xu, Xianfeng Gu, and Min Zhang. Autoencoder- based conditional optimal transport generative adversarial network for medical image generation. Visual Informatics, 8(1):15-25, 2024.
Zheyu Oliver Wang, Ricardo Baptista, Youssef Marzouk, Lars Ruthotto, and Deepanshu Verma. Efficient neural network approaches for conditional optimal transport with applications in bayesian inference. SIAM Journal on Scientific Computing, 47(4):C979-CI005) 2025.
Geoffrey Watson. Smooth regression analysis. Sankhyā: The Indian Journal of Statistics, Series A, pp. 359-372, 1964.
Ruofeng Wen, Kari Torkkola, Balakrishnan Narayanaswamy, and Dhruv Madeka. A multi-horizon quantile recurrent forecaster. star, 1050:28, 2018.
KD Yang and C Uhler. Scalable unbalanced optimal transport using generative adversarial networks. In International Conference on Learning Representations, 2019.
Qinqing Zheng, Matt Le, Neta Shaul, Yaron Lipman, Aditya Grover, and Ricky TQ Chen. Guided flows for generative modeling and decision making. CoRR, 2023.
Haoyi Zhou, Shanghang Zhang. Jieqi Peng. Shuai Zhang, Jianxin Li, Hui Xiong, and Wancai Zhang Informer: Beyond efficient transformer for long sequence time-series forecasting. Proceedings of the AAAI conference on artificial intelligence, volume 35, pp. 11 11106-11115, 2021.
Zhengbang Zhu, Hanye Zhao, Haoran He, Yichao Zhong, Shenyu Zhang, Haoquan Guo, Tingting Chen, and Weinan Zhang. Diffusion models for reinforcement learning: A survey, arXiv preprint arXiv:2311.01223 2023.
14

---

Published as a conference paper at ICLR 2026
A FURTHER APPLICATIONS OF HTI
We will now elaborate on some especially compelling potential applications for HTT.. In general, HTI can be useful in scenarios when user is deploying NN in dynamic environment, where behavioural preferences are context-dependent, and where the NN has a hyperparameter with a known, tangible behavioural effect. Traditionally, such deployment scenarios, user would either have to compromise with some fixed NN behavioural setting, determined at training time, or allow dynamic behaviours by undergoing slow and expensive retraining different hyperparameter settings when deemed necessary. HTT can alleviate this by enabling much faster inference time NN behavioural adaptation, by sampling estimated outcomes from the surrogate model p(y|r, J X) for a novel À setting. For visual depiction of HTI in action, see Figure 4
A.1
VARYING NEURAL NETWORK ROBUSTNESS IN DYNAMIC NOISE SETTINGS
Perturbations (e.g. Gaussian noise) of magnitude added to NN training data can increase robustness during inference for noisy inputs (Goodfellow al. 2015 Madry et al. 2018 Calibrating the training noise to that expected to be seen in deployment can lead to optimal results in terms of inference-time accuracy. The hyperparameter directly controls this trade-off: higher typically increases robustness to noisy inputs but may decrease accuracy on clean data.
Consider an image classification NN used quality control system on manufacturing line, where the input Timage is an image of a product. The desired level of robustness e" might change based on several factors:
Environmental conditions: Changes in factory lighting can alter image noise.
Operational mode: A user might decide to temporarily increase sensitivity to minor defects (requiring lower e for higher aceuracy on subtle features) during a specific batch run, or prioritise overall stability (higher e²) if the line is known to be experiencing vibrations.
Sensor age: As the camera ages, its noise profile might change, warranting an adjustment to c"
HTI would learn surrogate model P(yetass Timage: c). At inference time, based the current conditions and any explicit user preference for robustness, an appropriate e" can be selected. The system then samples from p[Velass 2image E") )") to obtain predictions if from model specifically tuned for that desired robustness level, without needing on-the-fly retraining.
A.2 VARYING SHORT- VS. LONG-TERM FOCUS IN REINFORCEMEN LEARNING
The discount factor 7 E [0, 1) in reinforcement learning (RL) determines an agent's preference for immediate versus future rewards. A low y leads to myopic, short-term reward-seeking behaviour, while a y closer to encourages far-sighted planning, valuing future rewards more highly.
Consider an RL agent managing patient's 's chronic disease treatment, such as Type 2 Diabetes, where actions involve adjusting medication dosage or recommending lifestyle interventions. The state includes physiological markers (e.g., blood glucose levels, HbAlc) and patient-reported outcomes. The optimal planning horizon, and thus the desired discount factor 7 can vary based on patient preference. For example, patient might express desire prioritize aggressive short-term glycemic control before an important impending event, or prefer more conservative approach at other times when they know their activity will be low. With HTI, users could then adjust the desired > based on the current clinical context. The system would then sample actions from p(a]s, 7*), allowing the treatment strategy dynamically shift its focus between immediate needs and long-term objectives without retraining the entire RL policy for each desired
A.3 VARYING FIDELITY AND DIVERSITY IN GENERATIVE MODELLING
Variational Autoencoders (VAEs) Kingma& & Welling 2014) are generative models that learn a latent representation of data. The 3-VAE (Higgins et al. 2017 introduces hyperparameter 3 that modifies the VAE objective function by weighting the Kullback- Leibler (KL) divergence term, which acts as a regulariser the latent space. The choice of B critically influences the model's behaviour:
15

---

typically involves Published as a conference paper at ICLR 2026
Mr 18 (-1281)
A. 0
Mao
A.
pors (-120)
(6 (-1 A))
A.
por, (-1201)
A.
po (-1201)
rels (:1270
(f(1|21, 2))ACA
A. 0.9
0.6
A. 0.1
P8x2 (-120)
R
Figure Example inference-time adjustment enabled by HTI We illustrate disparate user preferences affecting desired NN behaviour (desired À level) for different users in this abstract example. Having a fixed number of trained NNs (pex,) ) only allows partial exploration of the full hyperparameter trajectory, while an HTI surrogate model ((p((zz,, A))xeA) can estimate outputs across the entire spectrum of hyperparameters (estimated conditional probability paths represented by solid blue/red lines). Crucially, syyperparameterinnucce dynamics can differ amongst input conditions (x;)) as the true conditional distributions move along their own respective manifolds (Mz;), an effective HTT model must learn conditional dynamics.
Low B (e.g., B 1): With less pressure the KL divergence term, the model prioritizes reconstruction accuracy. This often leads to generated samples with high fidelity (i.e., they closely resemble the training data and are sharp/realistic). However, the latent space might be less structured or more "entangled," potentially leading to lower diversity in novel generations and poorer disentanglement of underlying factors of variation.
High 3 (e-g., 3 > 1): A higher 33 places more emphasis on making the learned latent distribution (=|1x) close to the prior p(=) (typically a standard Gaussian). This encourages a more disentangled latent space, where individual latent dimensions might correspond to distinct, interpretable factors of variation in the data (Burgess et al. 2018). While this can lead to greater diversity in generated samples and better generalisation for tasks like latent space interpolation, it might come at the cost of reconstruction fidelity, potentially resulting in blurrier or less detailed samples as the model sacrifices some reconstruction capacity to satisfy the stronger regularisation.
Consider 3-VAE trained to generate images. If user needs to generate photorealistic images, lower 3 would be preferred to maximise the sharpness and detail, ensuring the generated image is of high perceptual quality. On the other hand, if user is brainstorming image idcas, a higher 8 would be beneficial, encouraging the model to generate a wider variety of images and styles, even if individual samples slightly less photorealistic. HTI could learn surrogate generative model p(yimage x, 3). The user could then dynamically adjust 3* based on their current task.
A.4 FLEXIBLE HYPERPARAMETER OPTIMISATION WITH BAYESIAN OPTIMISATION
Standard Bayesian Optimization (BO) Snoek et al. 2012 Shahriari et al. 2015
learning a probabilistic surrogate model for a specific sealar objective function f R (e.g.. validation accuracy). This creates rigid dependency: if the user's preference changes during deployment-for instance, shifting from maximising pure accuracy to maximising accuracy subject to faimess constraint or inference latency budget- the learned surrogate is no longer valid for the new objective, and the hyperparameter search process must be restarted.
HTI can decouple the surrogate model from the objective function. Because HTI learns surrogate for po, (y/z) rather than a scalar objective, it can be used to calculate any performance metric derived from the model outputs
16

---

Published as a conference paper at ICLR 2026
1. The HTI model is trained on sparse set of anchor models to learn the conditional probability paths.
2. Post-training, user can define an arbitrary objective function [()) based on the model outputs (e.g., Expected Calibration Error, Fl-score, or custom utility function balancing risk and reward).
3. A BO optimiser searches for the optimal X that minimises J(A) by querying the HTI surrogate P(y|a, A).
Critically, evaluating the objective J via the HTT surrogate is much faster than retraining the original neural network. This could allow users to explore arbitrary Pareto frontiers of competing objectives without the need for further expensive ground-truth model training, or training multiple surrogate objectives as in standard BO.
17

---

Published as a conference paper at ICLR 2026
B SAMPLING ALGORITHM
We summarise our sampling procedure, as detailed in 44 below Algorithm 2
Algorithm 2 Sampling from p(y)
Require:
True distributions {Ptx (-1x))tTTT.. CLOT maps (Tor, h. geodesic path generator Sōg. target marginal t",, condition x E X
Find such that txs 6x+1 yv Pt. (-|x). 0k+1 = Tor, (Uk+2))
Tobs and tu < xt lk+1.
Define spline geodesic path (x(+) with ip = Sog (ykks Uk+1, x)
s" = (t" (t)/((x++ 12).
& Normalise target marginal for current interval
9r- = (s*)
return V/-
18

---

ot), which we Published as a conference paper at ICLR 2026
C EXPERIMENTAL DETAILS
We now provide detailed experimental set-ups for each of our experiments in 5 The code implementation of our method included here: https:://itth// com cow./aarraa2////// hyperpaarameteertraaaaaaa
inference
We ran all experiments between 5-10 minutes, depending on the surrogate model. To produce data for the wweih--eigg experiments, took 3.5 hours train PPO agent each setting. For the cancer experiment, it took between 2-15 minutes to train the surrogate models. For the reacher experiment, it took between 1-7 minutes train the surrogate models. For the quantile regression experiment, it took approximately 5 minutes to train the MLP quantile forecasters, and between 2-15 minutes to train the surrogate models. For the generative modelling dropout experiment, it took between 6-20 minutes to train the surrogate models. Ultimately, final experimental runs involve approximately 120 GPU hours of training time.
Azure VM A100 GPU. A single run for the semicircles experiment took
We base the implementation of our method off the code from Pooladian er al. (2024) (CC BY-NC 4.0 License, adapt for our specific setting.
C.I SEMICIRCLES EXPERIMENT
C.I.1 SEMICIRCLES DATASET
We describe here the temporal process used to generate the conditional semicircles synthetic dataset from $5.1 The dataset comprises points (a, y) associated with one of four discrete conditions {I. 2, 3, 4]. generated at continuous time t [0, 1]. For each condition and time, points are generated by first sampling an angle from Von Mises distribution (von Mises 1918), with time- and condition-dependee centred around a unit circle radius. Specifically, the generation process for single point under condition at time t is follows:
mean, and then sampling a radius from a Log-Normal distribution
Global parameters:
Tnom Nominal radius of the semicircles.
Frad 0.05: Standard deviation of the logarithm of the radial component, controlling radial
spread.
Kang 5.0:
Angular concentration parameter for the Von Mises distribution.
Generation:
For each condition e and time t:
1. Sample radius (R): The radial component R is drawn from a Log-Normal distribution, such that log( R) is normally distributed:
log( [) N Fog
where Mog log((nom) and Olog ardd Thus,
R &~ LogNormal( a
This distribution independent of condition e and time
2. Mean angle (Jany (c, f)) and semicircle centre (Fottet,e): The mean angle Mang (c, t) and the x x-coordinate of the semicircle's 's center Toffset, are determined by the condition c and time t:
1.0
if E (1,2] if E (3,4)
Toffistt.
1.0
lπ
if c = 1
(top-left semicircle, 0 + T)
Mang Mang(c, - tπ
if 2
(bottom-left semicircle, 0 -T)
(1 1)TT
if 3
(top-right semicircle, क -+00 0)
(t 1)TT
if = 4
(bottom-right semicircle, - π + 0)
19

---

Published as a conference paper at ICLR 2026
3. Sample angle (t)(t)): The angular component @c(t) is drawn from Von Mises distribution centred at the mean angle:
P (t) VonMises (c, t), Kang))
Cartesian coordinates (t), We (t)): The 2D coordinates are obtained by converting the sampled polar coordinates (R, (c))) to Cartesian, relative to the semicircle's ' center:
(t()) e + R (os())))))
(e(t) (in(
The full dataset a given time consists of N samples drawn from each of the four conditional distributions.
In $5.1 and 5.5 our training data consists of 100 samples for each condition at times E {0. 0.5, 1.0). The geodesics plotted in Figure begin at true points sampled at t = 0 and end at their estimated CLOT maps at For the numerical results in Tables [ and we compare estimated distributions from the respective models to the true distributions at E (0.25, 0.75).
C.1.2 MODEL DETAILS
The hyperparameters for the surrogate models used in the semicircles experiments are listed in Table Note that, since have discrete conditions in this experiment, construct separate NW density estimators for each condition, hence we set hy as N/A.
Hyperparameter
Value
0.05 for models with ú7 0 otherwise
hg
0.05
hr
N/A
Epochs
2001
Goa learning rate
10-3
Gea MLP hidden layer sizes
[128, 128] ReLU
Goa activations
Gea eigenvalue budget
2
g0, , Tor MLP hidden layer sizes
[64, 64, 64, 64] [1024, 1024]
Say MI.P hidden layer sizes
go,- Toy Sog learning rate
10-4
ge, Ter. Sgg activations
ReLU 15
Spline knots
FiLM layer size (applied to first layer activations)
16
T-Transform solver
LBFGS, 10 iterations
Min-max optimisation
1 x Go update per 10 x go,- Tez So updates
Table 8: Hyperparameters for semicircle experiments 5.1
C.2 CANCER THERAPY EXPERIMENT
C.2.1 ENVIRONMENT
We conduct this experiment using the "GhaffaiCannerEEnn-cotinnooo environment from DTR- TTR-Bench/DTR-GGG (Luo et al.) (024) (nttps://git http//gttttu..////////// cow/ Sttn///ileLLoooTRRR.. MIT license) which based on the mathematical model treatment of cancer with metasta- sis using radiotherapy and chemotherapy proposed in Ghaffari et al. (2016). The implementation deviates from Ghaffari et al. (2016) by treating the dynamics of circulating lymphocytes (c,) and tumor-infilrraiing eytotoxic lymphocytes (c2) constant.
The state at time an 8-dimensional continuous vector representing key biological and treatment- related quantities:
S = =Tpp,, Np,t Lep.rs Ct, T No. L,d, L,a,M]" M1)
where:
20

---

Published as a conference paper at ICLR 2026
T p.t: Total tumour cell population at the primary site.
ppr Concentration of Natural Killer (NK) cells at the primary site (cells/L..
pp Concentration of CD8+T cells the primary site (cells/L).
C Concentration of lymphoeytes in blood (cells/l.).
T., Total tumour cell population at the secondary (metastatic) site.
Ne.: Concentration of NK cells the secondary site (cells/L).
Ls,t: Concentration of CD8+T cells at the secondary site (cells/L.).
My: Concentration of chemotherapy agent in the blood (mg/L).
All state components are non-negative real values.
The action at time a 2-dimensional continuous vector representing the treatment intensities:
A = [De, v,]"
where:
Dr: The effect of radiotherapy applied at time t.
Vy: The effect of chemotherapy applied at time I.
These actions influence the dynamics of the state variables according to the underlying mathematical ODE model.
The reward Rg received after taking action A, in state S; and transitioning to state S is designed to encourage tumor reduction while penalizing significant deviations in Natural Killer (NK) cell populations, with additional reward penalty in terminal states. Let So [Tp.o. ........ be the initial state an episode. The components of the reward at each non-terminal step are:
Tumor reduction component (Rumor): This component measures the relative reduction in total tumor cells. First, the total tumor populations at the current step k (representing S2+1) and at initial step are calculated:
Tiot Tp + Ta,k
and Tin.0 Tp.o + Ts.0
These are then log transformed:
To = (mx Two,k)) and To In(mmx(ee, Two,o))
The tumor reduction reward then:
T To
Reumor, =
NK cell population penalty (Rak): This component penalizes deviations of the total NK cell population from initial value. The total NK cell populations are:
Num, Np, + Ns,k and Nww, Ng.0 + No.0
These are also log- -r-trannffrmee::
Ni In(max(e, Nun,))) and No In(max(e, Nia,o))
The penalty is then calculated, with weighting factor Anki
No+1
Rak, -
No
Finally, fermination reward (Riem) is added if the episode ends:
100
if positive termination (no more tumour)
R
100
if negative termination (max tumour size)
0
if non-terminal step
The total reward at step / is:
Ti-1 To
N
RS = Ratop. + Riem
+
Riem
No
21

---

Published as a conference paper at ICLR 2026
C.2.2 NON-LINEAR REWARD VARIANT:
For the non-linear reward scalarization experiment ($5...3) denoted as Cancer_nomiing modify the reward function to incorporate a hinge mechanism on the NK penalty term, employing the non-linear reward scalarization discussed eq. (6) in (Rädulescu al., 2020). In this setting, the weighted NK cell penalty is only active if the relative deviation exceeds threshold of 0.01. The modified penalty term Rax, is defined as:
-Ank
N No
M
0.01
R
No
0
otherwise
All other components of the reward function remain unchanged.
C.2.3 POLICIES
We train PPO agents (Schulman al. 2017) for the true distributions po, (als) at various Ank settings using the implementation in Stable Basel nes (Raffin et al., 202T) (MIT license, https://
ctm/ RM/ st able baselines3) with all other hyperparameters left at default, using the MLPPolicy architecture. For each agent, we train for 500, 000 timesteps.
Once trained, we use samples from the models with Ank E {0, 5, 10] as the training dataset for each surrogate model. Specifically, collecting 10 actions from each (stochastic) policy per observation. We evaluate each surrogate model Ank E 21, 4, 9.
the agent with Auk = 10 for 100 steps in the environment,
C.2.4 MODEL DETAILS
The hyperparameters for our surrogate models for the adaptive reward-weighting experiment are listed T
Hyperparameter
Value
0.01 for models with ú, 0 otherwise
hy
1.0
ha
1.0
Epochs
2001
Gog learning rate
5 10-3
Geo MLP hidden layer sizes
[128, 128)
Gea activations
ReLU
Geo eigenvalue budget
2
90,; Tor MLP hidden layer sizes
[64, 64, 64, 64)
Sa MLP hidden layer sizes
[1024, 1024]
90,- Tey Sog learning rate
10
90,- Toy, Sog activations
ReLU 15
Spline knots
FiLM layer size (applied first layer activations)
16
c- c-Transform solver
LBFGS, 3 iterations
Min-max optimisation
1 Ga update per 10 x 9u, Top, So updates
Table 9: Hyperparameters for our surrogate models in the cancer therapy experiment in 5.2.3
5.2. and
For the direct surrogate model, we train four-layer MLP using supervised learning, with inputs of the base action, condition, and target hyperparameter, and output the target action at the relevant hyperparameter setting. We list the direct surrogate hyperparameters in Table T0
For the CFM surrogate model (Lipman et al. 2023), we train two flow matching models, to model the
vector fields between the distributions Ank and Auk 5, and between Auk 5 and Ank 10
respectively. We base our implementation on the open source code from Lipman et al. (002). found here: https::// htgtttttbttttt httt://aaacoooooeeeeerc....c
httllooowwwwtoteeeee (CC BY-NC
22

---

Published as a conference paper at ICLR 2026
Hyperparameter
Value
Epochs
10000
Early stopping patience
100
Validation set
10%
Batch size
256
Learning rate
10-3
Hidden layer sizes
664, 64, 64, 64]
Activation function
Swish
FiL. layer size (applied to first layer activations)
16
Table 10: Hyperparameters for the direct surrogate model in the cancer therapy experiment in 5.2.1
4.0 License). We extend this implementation incorporate external conditions via FiLM layer. The hyperparameters for both of the CFM models in this surrogate model are listed in Table II
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] Swish
Activation function
FiL.M layer size (applied to first layer activations)
Table II: Hyperparameters for the CFM surrogate model in the cancer treatment experiment in 5.2.1
For the MFM surrogate model, we base our implementation on the open source code from Kapusniak et al. (2024), found here: https://github. com/ kkwpusniakuukkkkaak wwte-ffwwwwaachhteeeemmmmmmma-a- iigig (MIT License). This method is similar to flow matching, however aims to learn vector field that leads to interpolants staying the data manifold defined by the observed data. It does so by learning a NN-based correction to the straight line interpolants used associated with data-dependent
in CFM training, that designed to minimise the transport Riemannian metric. These corrected interpolants are then used to train neural vector field.
For the data-depended metric, we use their LAND formulation, which sets diagonal metric GLAND(z) = diag(h(z) + +e)) where
N (xf x 20
|| x;
ha (x)
expp(
(16)
20
with kernel size
We extend the original implementation to incorporate external conditions via a FiLM layer both the NN-based interpolant correction, and the neural vector field. The hyperparameters use for the MFM method are listed in [22
For the NLOT surrogate model, this is equivalent to method, but setting úí 0 and using the fixed eigenvalue metric from Pooladian et al. (2024). We use the same hyperparameters as our surrogate models for this method.
23

---

Published as a conference paper at ICLR 2026
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
Interpolant hidden layer sizes Vector field hidden layer sizes Activation function FiLM layer size (applied to first layer activations)
[64, 64, 64
664 64, 64
SELU
16 0.001
Table 12: Hyperparameters for the MFM surrogate model in the cancer treatment experiment in
5.2.1
C.3
REACHER
C.3.1
ENVIRONMENT
We conduct this experiment using the Reacher-v2 environment from OpenAI Gym (https 7/ thub. com/ openai/gym gym, MIT License). This environment consists of two-jointed robotic arm where the goal is to move the arm's end-effector to a randomly generated target location.
The state at time is --ddiennioonl continuous vector representing the angles and velocities of the arm's joints, as well as the location the target and the vector from the fingertip to the target:
S, (co)(O1), cos (00()0))) (in(
sin0)) (f2) Furget Marget, 0,, 02 Tlingentp
Turgett Vingertip
Marget Fingertip
Atargtt
where:
cos(f1). cos((2))) Cosine of the angles of the two joints.
sin(f1), sin 02): Sine of the angles of the two joints.
Turget: Varget The and y coordinates of the target location.
01, 0z: The angular velocities of the two joints.
Zfingertip Ttargets Yfingertip
Yurrgt Ffingettip
Staget: The vector from the fingertip the target.
The action at time t is 2-dimensional continuous vector representing the torque applied to the two joints:
A, = [7, T2)
and each E 1].
The reward RG received each step is the sum of a distance-to-target reward and a control cost penalty:
R = Pingeenip Purger ||2 Acomtral |i
where the first term is the negative Euclidean distance between the fingertip and the target, and the second the negative squared Euclidean norm of the action vector, which penalises large torques. We introduce the weighting hyperparameter, Aguntrol, that controls the strength of the control penalty
in the reward.
C.3.2 POLICIES
We train PPO agents Schulman et al. 2017) at the setting Acmmmol E 21, 3,5, using the im- plementation in Stable Base n (Raffin al. 2021) (MIT license, https: /e/g:://ttttt thub
24

---

Published as a conference paper at ICLR 2026
We the MLPPolicy architecture with default hy- perparameters. Each agent trained for 1, 000, 000 total timesteps.
Once trained, samples from the models with Acontol {1.5} the training dataset for surrogate model p(a)s s, A). Specifically, we run the agent with Acontrol = 1 for 1000 steps in the environment, collecting actions from each policy per observation. We evaluate each surrogate model
at Acontrll E (2,3.4).
C.3.3 MODEL DETAILS
The hyperparameters for our surrogate models for the Reacher experiment are listed in Table 13
Hyperparameter
Value
0.001 for models with 27, 0 otherwise
Q hy
2.0
ha
1.0
Epochs
2001
Go learning rate
5 10-3
Gea MLP hidden layer sizes
[128, 128]
Gsg activations
ReLU
Gea eigenvalue budget
2
90,; Tor MLP hidden layer sizes
[64, 64, 64, 64]
Sag MLP hidden layer sizes
[1024, 1024]
_ Toy Sog learning rate
10-4
ge, Te Seg activations
ReLU 15
Spline knots
FiLM layer size (applied first layer activations) o-Transform solver
16 LBFGS, iterations
Min-max optimisation
1 Geo update per 10 go,- Tey, So updates
Table 13: Hyperparameters for reacher experiment in 5.2.2
For the direct surrogate model, we train a four-layer MLP in the same fashion the cancer therapy experiment, with the same hyperparameters (Table 10).
For the CFM surrogate model, train flow matching model, between the distributions at
Aseatrol = 1 and Acontrol = with the same hyperparameters as in the cancer experiment (Table IT
For the MFM surrogate model, we train with the same hyperparameters in the cancer experiment (Table 12), but with = 1.
C.4 QUANTILE REGRESSION
C.4.1 DATA
We the ETTm2 dataset from the Electricity Transformer Temperature (ETT) collection (Zhou
et al. 202T) (https: /tt::://gttttt com/ / CC BY-ND 4.0 License), which contains data on electricity load and oil temperature. We formulate forecasting task for oil temperature, with an input horizon of 12 steps to predict an output horizon of 3 steps. The dataset partitioned chronologically, with the first 70% used for training the ground-truth models and the subsequent 15% for validation. From the remaining data, the next 1200 samples form the training set for the HTI surrogates, and the final 180 samples are used as the HTI testing set evaluate surrogate model performance.
C.4.2 GROUND- TROUT FORECASTERS
The ground-trune truth forecasters are three-layer MLPs with hidden dimensions of [256, 128, 128). We train separate model for each target quantile {0.01, 0. 0.25, 0.5, 0.75, 0.9, 0.99). Training is performed for up to 2000 epuchs using the pinball loss function, with a learning rate of 10 - and a batch size of 32. We employ early stopping with patience of 10 epochs.
25

---

Published as a conference paper at ICLR 2026
The pinball loss, L, 9), for true value and quantile forecast quantile level is defined as:
(y (v -
y y if y
This loss function penalizes under-prediction and over-prediction asymmetrically, which encourages the model to learn the specified quantile.
To create the HTT training dataset, T == 0.99 to generate forecasts on the 1200 inputs of the HTI training set. For evaluation, the forecasts from the remaining ground-truth models (for T E {0.1, 0.25, 0.5, 0.75, 0.9}) on the 180 HTI test inputs serve as the ground-truth quantiles.
the ground-truth forecasters trained for 0.01 and
C.4.3 MODEL DETAILS
The hyperparameters for our surrogate models for the quantile regression experiment are listed in Table
Hyperparameter
Value
0.01 for models with 24, 0 otherwise
Q by
1.0
h
1.0
Epochs
1001
Goo learning rate
5 10-3
Go MLP hidden layer sizes
[128, 128] ReLU
Goa activations
Goe eigenvalue budget
3
go Tor MLP hidden layer sizes
64, 64, 64, 64, 64, 64, 64, 64]
Sa MLP hidden layer sizes
[1024, 1024]
ge,- Ter Ses learning rate
10-4
90,- To,, So activations
ReLU
Spline knots
15
FiLM layer size (applied to first layer activations)
16
c-Transfomm solver
LBFGS, 10 iterations
Min-max optimisation
1 Geg update per 10 90,, To,, So updates
Table 14: Hyperparameters for ETT experiment in 5.3
For the direct surrogate model, we train an eight-layer MLP with hidden dimension of 64, to match the increase the number layers for the Kantorovich potential and CL.OT map MI.s in our surrogate models. The other hyperparameters
the same as in the cancer experiment (Table I0).
For CFM surrogate model, we also use an eight-layer MLP with a hidden dimension of 64 for flow matching model. The other hyperparameters are the same as in the cancer experiment (Table
For the MFM surrogate model, we train with the same hyperparameters in the cancer experiment (Table 12). but with o = 0.01.
C.5 DROPOUT EXPERIMENT
C.5.1 DATA
The two-moons dataset was generated with sklearn. datasets. www.maakemoommo using 2000 points and noise of 0.05. The 2D co-ordinates were standardised to zero mean and unit variance prior any modelling; the binary class label 0, was used as the conditioning variable. For the HTI training set we trained ground -ruth diffusion models at anchor dropout levels E {0, 0.5, 0.99) and drew 1000 samples from each converged model (balanced across classes).
26

---

Published as a conference paper at ICLR 2026
C.5.2 DIFFUSION MODEL
Ground-truth models were implemented DDPMs (Ho et al., 2020) with conditional MLP score network. The network receives the noisy coordinates E IR- a normalised timestep scalar / and the class condition 2, and predicts the noise R². The MLP has two hidden layers of size 256, ReLU activations and dropout applied after each hidden layer. The forward schedule uses T = 100 timesteps with 3 linearly spaced from le-4 to 0.02. Models were trained with Adam (learning rate le-3), batch size 128, and for 250 epochs. Sampling was performed via the standard ancestral reverse diffusion loop using the trained network and the same noise schedule.
C.5.3 MODEL DETAILS
Surrogate HTI models were trained to interpolate between the three anchor diffusion marginals and evaluated at held-out dropout settings E {0.1, 0.3, 0.4, 0.6, 0.7, 0.8, 0.9}. Evaluation compares the Wasserstein distance between surrogate pushforward samples and samples from ground-trute diffusion models trained at the target y values. Hyperparameters used to train our surrogates are listed Table [
Hyperparameter
Value
0.01 for models with úi 0 otherwise
hy
0.2
h
1.0
Epochs
2001
Gog learning rate
5 10-3
Gea MLP hidden layer sizes
[128, 128] ReLU
Geg activations
Geo eigenvalue budget
3
go, Tor MLP hidden layer sizes
[64, 64, 64, 64] [1024, 1024]
Sag MI.P hidden layer sizes
90,- Tor Sos learning rate
10-4
90,- Tey Sag activations
ReLU 15
Spline knots
FiLM layer size (applied first layer activations) -- TTransform solver
16
LBFGS, 10 iterations
Min-max optimisation
Gog update per 10 x go,- Top Sog updates
Table 15: Hyperparameters for dropout experiment in 5.4
For the direct surrogate model, we train a four-layer MLP in the same fashion as the cancer therapy experiment, with the same hyperparameters (Table T0).
For the CFM surrogate model with the same hyperparameters as in the cancer experiment (Table
For the MFM surrogate model, we train with the same hyperparameters as in the cancer experiment (Table 12), but with 0.5.
27

---

Published as a conference paper at ICLR 2026
D EXTENDING TO MULTIPLE HYPERPARAMETERS
limitation of our current approach is its design only immediately appropriate for a single, continuous hyperparameter. We see extensions multiple and discrete hyperparameter settings key direction for future research. Extending our current HTI method to this setting is non-trivial.
One simple extension that would allow for interpolation between multiple hyperparameters with current method involves extablishing mapping from the multi-dimensione hyperparameter space to single 'time' space, allowing our interpolation scheme that works on a single dimensional 'time' variable to apply. We have considered two representative strategies for creating such a mapping data-driven Principal Curve and geometric space-filling Hilbert Curve but there are outstanding limitations to both potential approaches.
Principal Curves: A principal curve, non-linear generalisation of PCA, is the smooth curve that captures the most variance a dataset. If we have multiple observed sittidiimmniinall hyperparameters, could find the principle curve through them, which could serve as ID "time" axis. The primary limitation of this approach is that it only allows for interpolation hyperparameter settings defined along this learned curve. To approximate arbitrary setting that is not on the curve, one would first have to project it onto the curve.
Conversely, a space-filling Hilbert Curve is a pre-defined geometric con- struction whose single, continuous line is guaranteed to pass through every point in multi-dimensional space, ensuring full coverage. While this could guarantee coverage, its critical flaw is that breaks locality. Our method is grounded in Optimal Transport and least-action principles, which assume that small changes in our "time" variable should lead to small changes in the output distribution. Hilbert curve would not necessarily respect this intuition, potentially mapping two very different distributions to be 'temporal neighbours'.
Hilbert Curves:
28

---

Published as a conference paper at ICLR 2026
E SPARSITY INVESTIGATION
To investigate sensitivity of different surrogates to data sparsity, we evaluate HTTI performance with various number of anchor distributions in the Cancer and Reacher environments. We range from the sparse settings from 1 and :[5.2.2 to a dense setting, where training data is available at every evaluation setting. Figure 5 shows that, in both environments, the performance gap between methods is negligible in the dense regime, where interpolation is trivial, and this widens as sparsity increases. Our method degrades the least as interpolation becomes more difficult, confirming the effect of our inductive biases and their importance in sparse settings.
120
-56
80
Reward
Direct CFM
Direct
CFM
20
62
-64
-20
K
R 1-
R4
-68
11
2
Number of Anchors
Number of Anchors
Figure 5: Surrogate model reward in Cancer (left) and Reacher (right).
29

---

Published as a conference paper at ICLR 2026
F LLM USAGE
In this work, we used LLMs to assist with the writing of this manuscript. This primarily involved consulting LLMs refine drafts, improving the coherence and clarity of our work, and simplifying the writing process.
We also used LLMs during the experimental process, to help write code.
30


Reviews:
[{"review_id": "fjFkhFzy4d", "reviewer": "Reviewer_9CGQ", "summary": "The paper defines Hyperparameter Trajectory Inference (HTI): given a neural model whose behaviour changes with a continuous hyperparameter, learn a surrogate that interpolates the model’s output distributions between a few trained anchor values. To do this, the authors extend neural Lagrangian optimal transport to the conditional setting: they learn, from sparse observations, a conditional Lagrangian and amortised CLOT maps/geodesics so that inferred paths stay low-action and remain in dense regions. They validate this on a 2D toy task, on two RL surrogates (cancer therapy, Reacher), and on quantile regression, where their full model (learned metric + density potential) is consistently better than ablations and the CFM/direct surrogates.", "questions": "- Could you explain why exactly you need to introduce the $c$-transform instead of the standard dual formulation? \n- Concerning the move from classical COT to Lagrangian COT, could you articulate more clearly why a path-based cost is needed for CTI/HTI and whether using only the kinetic term with a learned metric (no potential) would already be enough for sparsity?\n- How does Equation 12 help in evaluating Equation 11 faster? As you train $T\\_{\\theta}\\_{T, k}$ via $\\mathcal{L}\\_{map}$, do you explicitly need to minimize Equation 11?\n- Would you expect to obtain so clean results on RL experiments if the reward scalarization function were non-linear?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "- HTI is a useful abstraction: many hyperparameters in practice (reward weights, discount factors, quantile levels, robustness coefficients) define families of policies or predictors, but current practice trains only a few points on that curve. Framing this as conditional trajectory inference and tying it to OT provides a principled way to discuss “hyperparameter-induced dynamics.”\n- The authors do not simply make their approach learn pairwise conditional OT maps. It learns the cost itself via a conditional Lagrangian, with two inductive biases, and the Lagrangian is the device that lets them force the inferred paths to stay plausible when data are sparse.\n- The experiments section seems to confirm that the approach enables learning from a few, sparse observations...", "weaknesses": "- ... but all (few) real evaluation tasks are fairly low-dimensional on the output side, forecasts with short horizons. The method is sold as applicable to “complex and higher-dimensional geometries,” but the experiments do not necessarily convey this statement.\n- The RL scenarios considered are in fact well-behaved; the parameter $\\lambda$ yields a linear combination of \"objectives\" (e.g., main reward+penalty/cost). In such cases, it is widely known that the resulting trade-offs from tuning $\\lambda$ are well-behaved (see [1]).\n- The introduction of core mathematical notions (conditional OT, the c-transform, and the Lagrangian formulation) is difficult to follow for readers not already fluent in OT theory and Lagrangian. For instance, the paper directly presents the dual form with the c-transform (eqs. 2-3) and the action integral (eq. 5) without first motivating them intuitively or connecting them to the simpler Kantorovich dual with explicit constraints. As a result, the rationale for moving from standard OT to a Lagrangian formulation, and why this path-based cost is required for CTI/HTI, remains implicit.\n- The same issue carries over when explaining why the method can generalize from sparse observations: the text hints that the least-action and density biases provide this capability, but the link is scattered across sections 3 to 5 instead of being clearly stated when these concepts are first introduced. Overall, the exposition of the theoretical background is technically correct but conceptually opaque, making the narrative harder to follow than necessary.\n- As acknowledged by the authors, a main limitation of the approach is that it only scales to a unique hyperparameter. In practice, many hyperparameters must be tuned. This fact is acknowledged only at the end of the paper; the limitation should be explicitly stated earlier in the text.\n- Since one of the key selling points is the capability to learn from sparse observations, the evaluation would benefit from evaluating the quality of the predictions with respect to the sparsity of the observations (add or remove anchors and analyze the impact).\n\n[1] Roxana Radulescu, Patrick Mannion, Diederik M. Roijers, Ann Nowé: Multi-objective multi-agent decision making: a utility-based analysis and survey. Auton. Agents Multi Agent Syst. 34(1): 10 (2020)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "zSyUHdtQys", "reviewer": "Reviewer_rFXY", "summary": "This paper studies the problem of hyperparameter optimization for training neural networks. To do this, the authors develop a conditional Lagrangian optimal transport framework that allows them to learn a Lagrangian for problems of interest. The authors demonstrate that the learned Lagrangian can accurately model hyperparameter landscapes across examples, demonstrating the utility of their method.", "questions": "- Rather than just empirical evidence across examples, can the authors give a theoretically motivated way to choose which method to use? In other words, how do we guide our choice of conditional Lagrangian? Is there an automatic way to do this?\n- Does the method do a good job of extrapolating outside of the range of learned hyperparameters? Having a clear discussion of interpolation versus extrapolation in trajectory inference like this is useful.\n- I cannot get a sense of how scalable this method is. How expensive is this to run in really large scale ML settings?\n- Can the authors combine their work with Bayesian optimization formulations to yield a way to choose optimal hyperparameters rather than just modeling across a range of hyperparameters?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "- The paper does a good job of setting up its methodology, with clear explanations of each piece. \n- The problem formulation developed here seems new and is well motivated.\n- The use of learning a conditional Lagrangian Optimal Transport model is novel and elegant.\n- The authors give a range of experiments that demonstrate the ability of this method to learn hyperparameter trajectories.", "weaknesses": "- It is hard to see what the real use of this method will be, and the message of this gets somewhat lost in the experiments.  While this is a cool idea, what are the applications in which this will be most relevant, especially when one doesn't know a priori how chaotic the trajectories will be?\n- There are no baselines that this work is able to compare against.\n- While the results indicate that the method can do a good job of modeling trajectories with proper tuning, I am missing what insights we gain from these trajectories. Rather than just chasing performance, can the authors comment on interesting findings in what is actually learned?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "QxfFlai25n", "reviewer": "Reviewer_Ei6o", "summary": "This paper introduces a technique for performing hyper-parameter trajectory inference (HTI), a problem where one introduces a model $\\hat{p}(y|x,\\lambda) \\approx p\\_{\\theta\\_{\\lambda}}(y|x)$, which eases the adjustment of $\\lambda$ at inference time. To that end, the authors propose using Lagrangian Optimal Transport (Villani, 2009; Pooladian et al., 2024) to build their model $\\hat{p}$.", "questions": "The hyper-parameter trajectory notion seems a bit contradictory to me. For instance, usually hyper-parameters are fixed during training, so they don't really evolve. From what I got from the text, the authors are trying to model the mapping $\\lambda \\mapsto p\\_{\\theta\\_{\\lambda}}$ (see the intro) but in my view this lacks a temporal structure, i.e., $\\lambda(t), t \\in [0 , 1]$ to warrant the nomenclature \"trajectory inference\". Could the authors elaborate on this idea?", "limitations": "", "rating": 2, "confidence": 3, "soundness": 2, "presentation": 1, "contribution": 2, "strengths": "There are 2 main strong points with this submission.\n\nFirst, the authors do a nice link in their experiments section, with their motivations in the introduction. The reasoning for why one needs hyper-parameter inference are also clear and well motivated.\n\nSecond, there is a good variety of important tasks for which the proposed method applies. In all tasks the authors show an improvement over other methods, sometimes with a significant margin.", "weaknesses": "Overall, I think this paper has major organization issues:\n\n1) Section 4 is way too short to merit being a section on its own (2 paragraphs), and I think the authors define the problem they are treating way too late in the paper. In my view, this should be done as early as possible (e.g., in the beginning preliminaries section).\n\n2) With `1)` in mind, the preliminaries section seem disconnected with the problem statement. It could be nice to highlight how conditional OT/lagrangian OT relate to the HTI task.\n\n3) The paper seems to have many interesctions with the work of (Pooladian et al., 2024). For instance, they use the same amortization procedure and metric parametrisation. With that in mind, it would be nice to have a comparison with their method, highlighting the improvements over previous art.\n\n4) About section 5 in general, the fact that the main algorithm is only shown in the appendix harms readability. Furthermore, the authors don't present the final objective function/learning problem being optimized.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)