Test Time Scaling for Neural Processes
Authors: Hyungi Lee
Moonseok Choi
Hyunsu Kim
Kyunghyun Cho
Rajesh Ranganath
Juho Lee
Venue: NeurIPS 2025 poster ()
Keywords: Neural Processes
Sequential Monte Carlo
Test Time Scaling
Latent Variational Posterior
Primary area: probabilistic_methods

Abstract: Uncertainty-aware meta-learning aims not only for rapid adaptation to new tasks but also for reliable uncertainty estimation under limited supervision. Neural Processes (NPs) offer a flexible solution by learning implicit stochastic processes directly from data, often using a global latent variable to capture functional uncertainty. However, we empirically find that variational posteriors for this global latent variable are frequently miscalibrated, limiting both predictive accuracy and the reliability of uncertainty estimates. To address this issue, we propose Test Time Scaling for Neural Processes (TTSNPs), a sequential inference framework based on Sequential Monte Carlo Sampler (SMCS) that refines latent samples at test time without modifying the pre-trained NP model. TTSNPs iteratively transform variational samples into better approximations of the true posterior using neural transition kernels, significantly improving both prediction quality and uncertainty calibration. This makes NPs more robust and trustworthy, extending applicability to various scenarios requiring well-calibrated uncertainty estimates.

Test Time Scaling for Neural Processes
Hyungi Lee Kookmin University
Moonseok Choi KAIST
1hk270@@kookmin..e.
.r.
Hyunsu Kim KAIST
Kyunghyun Cho
New York University&Geenetcch
kim. hywssuskaist..cc.k
kyunghyun. cho@nyu. edu
Rajesh Ranganath
Juho Lee KAIST
New York University
rajeshrécims..
nyu. www.mmmsnyueeee..
juholee@kiist.cm
Abstract
Uncertainty-awan meta-learning aims only for rapid adaptation to new tasks but also for reliable uncertainty estimation under limited supervision. Neural Pro- cesses (NPs) offer flexible solution by learning implicit stochastic processes directly from data, often using a global latent variable to capture functional un- certainty. However, we empirically find that variational posteriors for this global latent variable are frequently miscalibrated, limiting both predictive accuracy and the reliability of uncertainty estimates. To address this issue, we propose Test Time Scaling for Neural Processes (TTSNPs), a sequential inference framework based on Sequential Monte Carlo Sampler (SMCS) that refines latent samples test time without modifying the pre-trained NP model. TTSNPs iteratively trans- form variational samples into better approximations neural transition kernels, significantly improving both prediction quality and un- certainty calibration. This makes NPs more robust and trustworthy, extending applicability to various scenarios requiring well-calibrated uncertainty estimates.
the true posterior using
Introduction
In many real-world scenarios, tasks naturally vary in several aspects, such the number of avail- able data points and the degree of correlation between training and evaluation data. Despite the diversity, common challenge remains: each task typically offers only limited supervision, which makes it difficult to train models that generalize reliably within individual tasks [2, 27]. powerful paradigm. While meta- learning [15, 25, 22] has demonstrated significant success in improving adaptability in various fields via extracting transferable knowledge, reliable uncertainty estimation along with accurate predic- tions becomes important, as it allows models to express appropriate levels of confidence in their pre- dictions and make trustworthy decisions given limited information. To this end, uncertainty-aware meta-learning [48, 13, 49] has been proposed to integrate uncertainty quantification directly into meta-training and evaluation processes, promoting both rapid adaptation and calibrated prediction.
To address this issue, meta-learning [50, 15] has emerged
Among various frameworks, Neural Processes INPs; 18, 19] have emerged as a flexible and data- driven approach for ancertainty-aware meta-learning. By training parametric neural networks to maximize predictive likelihoods across meta-task datasets, NPs can capture the data-generating mechanisms underlying diverse tasks without relying on explicit prior assumptions. As result,
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

they offer a powerful balance between flexibility and uncertainty estimation, positioning them as foundational model for uncertainty- aware meta-learning [33]. common approach to modeling un- certainty in NPs involves introducing a global latent variable, which encodes functional uncertainty or epistemic uncertainty across tasks.
However, in this work, we empirically demonstrate that previous approaches on NP mod- els [19, 28, 33] leveraging global latent variable model functional uncertainty often suffer from miscalibrated variational posteriors for the global latent variable. Moreover, we find that such mis- calibration issues consistently arise regardless of the specific loss functions or model structure used during the training of NPs. This miscalibration presents a fundamental challenge for NPs, where the global latent variable serves as the core mechanism for conveying contextual information to target predictions. When the variational posterior fails to faithfully reflect the uncertainty in the context set, the resulting latent representations may become either overly deterministic or excessively diffuse. Consequently, the model's predictions on target points can suffer from poor calibration and reduced expressiveness, particularly when the target inputs deviate from those observed in the context set. These limitations undermine the generalization ability and reliability of NPs, especially in scenarios where accurate uncertainty estimation is crucial.
To overcome this limitation, we propose principled sequential inference framework based on Se- quential Monte Carlo Samplers (SMCS) [11, 10], designed refine latent variable samples at test time without modifying the pre-trained NP model. Specifically, our approach, Test Time Scaling for Neural Processes (TTSNPs), iteratively transforms latent samples from the variational posterior into more accurate approximations of the true posterior using SMCS, with the following two key components: (1) learned neural network transition kernels, and (2) learned intermediate distribu- tions, achieved by generating pseudo context points via a neural network. This refinement improves predictive performance and enhances uncertainty calibration, enabling NPs to produce more reliable and trustworthy predictions. As result, TTSNPs significantly extend the applicability of NPs to various tasks that demand robust and well-calibrated uncertainty estimation.
2 Preliminaries
2.1 Problem Setup
In typical NP settings, the dimensionality of datasets is assumed to be fixed [18, 19, 28, 32). In contrast, following the direction explored in Lee et al. [33], we consider more general setting in which the data for each task may vary in dimensionality, Specifically, we assume potentially infinite collection of tasks (Ty (Ti]JEN, each drawn i.i.d. from a task distribution pr(r). Every task Tj uses an input space (2)()) S RR==() and output space Vd,,o) S Rd)((), where dz, dy: N - specify each task's 's input-output dimensionality. Concretely, provides a dataset D, = (d 7i in which each dj.A (xj,e, v,, belongs to X2z Va,,,) A subset Cj & [n,] of indices defines the context set Dj.-- == [dj, kee,. while the complement 15 = [ny] \ Gj gives the target set Dj,t- Each dataset DD is considered i.i.d. from some unknown function fj: Xd=()) - Vo,(9)- Meta-learning aims to learn shared representation capturing how Xj.k and j,A related, using Dj,e (training) and Dj, (validation) across many tasks.
2.2 Neural Processes
NPs address the meta-learning problem by defining a distribution over functions f, conditioned on each context set Dj.c. The predictive distribution can be viewed as follows [33]:
p X,, = [II p(y fz, xxxx (1,)) df,, kE[n,]
(1)
where X, and Y, collect all inputs } and outputs } When adopting Gaussian likelihood for p(yl ff x) and introducing a latent variable ry E Rdd that parameterizes fy. one might write
p pY, (ic) = / II
Myya Mry (xxx)) or (xx,+))) ((r,) DD () dr,. kE(m))
(2)

---

where Φ is the parameter set of the NPs encoder module. NP variants differ primarily in how they encode Di. into rj. Conditional NPs (CNPs) [18, 21, 48] use a deterministic encoder that maps j.c to point estimate rj- Latent NPs [19, 17, 32] instead model uncertainty by placing a variational Gaussian over ry. Both approaches generate predictive means and variances via decoder.
(Krys o,,) = Saor (x,u, Fyi bb),
(3)
where 2b is the parameter set the NPs decoder module. CNPs maximize the predictive log- likelihood over meta-training tasks, while latent NPs optimize the predictive log-likelihood or an Evidence Lower BOund (ELBO) to balance data and Kullback-Leibler (KL) regularization. Here, ELBO is usually approximated by the following form:
Br, plogp(Y, DycYl 2 Bry
Σ kE|n,]
Egtr, D3) [log.N.a] - KI[g(r, ( || g((r P3.
(4)
where Ni. is shorthand for Myyj,. Mr, (xx,k), diag(o? (x).)))). In this paper, we are focusing on latent NPs, where we perform test-time scaling to sample from the latent variable more accurately.
2.3 Sequential Monte Carlo Samplers
In this section, briefly review Sequential Monte Carlo Samplers (SMCS) [11], which form the foundational framework for our test-time scaling approach. For a more comprehensive overview, refer the reader to Del Moral et al. [11], Dai al. [10].
Let 7 (xx) == / denote target distribution, where the unnormalized density 7(x) known pointwise but the normalization constant Z is unknown. SMCS aims to approximate expectations with respect to म, such as [r[ for a test function f. while simultaneously providing an estimate of Z. Rather than sampling directly from the target SMCS constructs sequence of intermediate distributions {= beginning with a tractable initial distribution To and gradually transporting These intermediate distributions are designed to interpolate between mo
toward the target TT and т, typically becoming increasingly complex as increases. Since the intermediate distributions generally not analytically tractable, SMCS employs sequential importance sampling adjust the samples accordingly.
.,..., In detail, let "/ be the unnormalized density of TT. i.e., (T( (<(x//Z for L T (note that Zo = as mo is assumed to be fully known). SMCS introduce forward transition kernels (Fr ( which propagates sample #T-1 to zt. Starting with zo mo, sample path To:T = (No -......... FT) is generated from the joint proposal with density,
T Q(xo.) = (=))) II Fe Iz-1).
(5)
In principle, one would like to use the marginal ʃ ( Jdzor- as a proposal density for n, but this is generally intractable. Instead, SMCS augments the target the path space 20:T with sequence of backward transitions [B.- xx) xt)) and defines the unnormalized path density
I((O.T) = I( Z where r(xor) = (xT) II
T
B B
(6)
t=1
sample path ogrr drawn from Q is then corrected with the importance weights,
wT(For) = [( (( B | ((0.)) To( IT= Fi IT-1)
Given the forward and backward kernels, SMCS starts by drawing N i..d particles [ from Tp. At each iteration, the particles are sequentially extended by drawing samples from the forward kernel F,, and the extended particles are re-weighted through the importance weights, which are computed in recursive fashion follows:
Ne (x B, (xi x1) for i 1, 1,...,A N.
mo(==)) 1, wx(b)))) wi- (o
(8)
7-1 [ 1) F((xt x}-r)
3

---

After completing the sequential update up to the step T, we have an estimator for the expectation E [f] and the normalization constant Z.
N
WT (*):)) Zj=1 wT(
E1 =
f(+)),
Z ~0
N Σ
wr (2).))
N
Learning for SMCS. Since both the intermediate distributions and the forward/backward kernels are crucial design choices in SMCS, it is beneficial to learn them when possible. Learning for SMCS can be formulated as a variational inference problem, where the ELBO provides a lower bound on the marginal likelihood defined by the SMCS procedure. A common approach is to begin with un- adjusted Langevin diffusions and modify the drift term (involving the score functions) to maximize the ELBO [12, 20, 53, 7]. Among these, Chen et al. [7] introduced an optimal control framework for tuning a continuous-timg extension of SMCS with resampling, which forms the foundation of proposed algorithm. Refer to Appendix B to additional explanations for SMCS.
3 Test Time Scaling for Neural Processes
3.1 Observation
As discussed in $ 2.2, latent NP models are typically trained by maximizing the predictive log-likelihood or ELBO. Ideally, this results in variational pos- terior (r))).) that approximates the true posterior ((r( x II(='')")ED p(y"|=' (1)"(=',)))p())) However, due to training limitations and overfitting to specific context distributions, the learned posterior often de- viates from the true posterior, especially at test time.
Libellbood
0.35
0.30
0.25
SMES
0.20
IMC
To validate this issue, we conducted experiments using simple NP model [19] trained on ID GP data with an RBF kernel. The same trained model was evaluated on unseen test datasets sampled un- der identical conditions. We compared the tar- get marginal likelihoods using latent samples drawn from the variational posterior, as well as from im- portance sampling [IS: 52], Hamiltonian Monte Carlo [HMC; 46], and SMCS [11] with Unadjusted Langevin Algorithm (ULA) transitions. Since all methods use the same encoder and decoder, differences in log-likelihood stem solely from the quality of latent samples. Fig. shows that the variational posterior underperforms compared to exact inference methods, rapidly saturating after 15-25 samples. This highlights its lack of diversity and poor calibration, indicating failure to capture epistemic uncertainty. Based these find- ings, we introduce a test-time scaling method using SMCS in the next section. We choose SMCS over other samplers due to its strong sample efficiency and fast convergence in high-dimensional regimes [42, 36, 62], which is also reflected in our empirical results.
15
23
33
30
100
200
1000
Number Samples
Figure 1: Marginal log likelihood compar- ison across variational posterior (VP), IS, HMC, and SMCS using the NP model [19]. The x-axis indicates the number of latent samples used for estimation.
3.2 Test Time Scaling with Sequential Monte Carlo Samplers
As mentioned in $ 3.1, even after being trained on large-scale datasets, the latent encoder in latent NP models often fails to fully capture functional eeceraiitty-paticullal when faced with struc- tured catoof-distribution test data, such as those arising from unseen kernels. To mitigate this issue, we propose a test-time scaling approach that enhances uncertainty estimation without requir- ing any modification to the model's pre-trained parameters. This design choice allows us to retain the representational power learned from the large-scale training data while refining the posterior inference process during evaluation.
Our key idea is to resample the global latent variable r, which governs functional uncertainty in latent NP variants. Since the quality of uncertainty estimation is directly influenced by how well this latent distribution aligns with the test context, we aim to improve via a principled sampling- based method. Specifically, we adopt the SMCS framework (11, 7] to transform an initial set of

---

trainedd lasent Encoder
learned
Pre trained Decoder
trained Encoder
X
Pierudo
SPC staps
Remorapligg art
GHC Stape
Target 1+?
Figure 2: Schematic for TTSNP. TTSNP refines the latent variables from the variational posterior into more accurate samples from the true posterior using the learned transition kernels.
samples drawn from the variational distribution q(r Dz) into set that better approximates the true posterior: To implement this, we first construct a sequence of intermediate distributions that gradually interpolate between the initial variational posterior and the target posterior. Then, particle transitions are guided by a parameterized forward kernel and backward kernel, both modeled using neural network. The forward kernel nudges samples toward high-probabillty regions of the target distribution, while the backward kemel helps estimate gradients and improve learning stability dur- ing training. We also define an SMCS- based training objective derived from KL divergence [41, 7] and the marginal oggllikelhood, which facilitates the learning of efficient and effective transitions for posterior refinement. To define an SMCS procedure, we must specify the initial distribution, intermediate distributions, and transition kernels. We detail our choices for each below.
Initial distribution. While the initial distribution is typically chosen be simple, in case with a pre-trained NP model, the variational posterior q(r Dz) = =N(r o serves as an effective starting point. Although may fully capture the uncertainty of the true posterior, provides an informative and computationally tractable initialization that can be further refined through SMCS,
Intermediate distributions. We construct intermediate distributions via a geometric annealing path, enriched with learned component. Specifically, for = -..... , T, we define
(r) p(r (2)) (x - q(r D. D 0
1 (10)
Here, the first term p(r DD.) is the true posterior. The second term includes recalibrated variational posterior 7(r D.) == Ar ,q 02), retaining the original mean fl but using fixed variance o². Empiricalll setting o² provides stable and robust performance. Together, these two factors form slight modification of the standard geometric annealing path without requiring additional learning. The third factor introduces our core contribution, where the additional factor is defined as a variational posterior conditioning on the context appended with pseudo context set 'D,-a synthetic future context predicted from the current estimate. Concretely, neural network h(D,,e) E) maps the original context D. and noise p(E) to pseudo context D,,- defining an implicit generative model for the future data. Our design is inspired by Martingale pos- teriors [16, 32]. where the predicted future data drive posterior uncertainty. In context serves as form "Jookahead", encouraging the model to maintain hypotheses that remain robust under plausible extensions of the observed context. This mechanism mitigates the overcon- fidence often observed in amortized inference by injecting uncertainty from possible yet-unseen data. Analogous "chain-of-Fe offtthought reasoning in language models, the pseudo context enables the model to simulate and prepare for potential future scenarios, thereby enriching the intermediate distributions with anticipatory structure and improving posterior coverage. For implementation, adopt the induced self-attention block [35] for h.
germutation-invarian
the pseudo
Forward and backward kernels. Given the intermediate distribution based on pseudo contexts, we define forward and backward kernels based on ULA. Following Chen et al. [7], we formulate our SMCS procedure as time-discretizations of the continuous-time dynamics governed by the following Stochastic Differential Equation (SDE) evolving over 7 [0, 1]. dr urr" r)dr f + odW,, r
(1I)
where u dd x [0, 1] + d is drift, is diffusion coefficient, and W,, is standard Brownian motion. parallel, we define the reverse-time SDE,
ds v(s", r)dr t odW,, si
(12)
5

---

with a drift v, diffusion and ,W, denoting time-reversed Brownian motion. We expect these two SDES be time-reversals of each other, and their marginals T₁ == t/T for ...... match the intermediate distributions we set, that is.
P = Pr Te for t O....,
(13)
where Pr, and IP, denote time- marginals of the forward and reverse SDEs at time TG. There potentially infinitely many choices of the pair (u, v), but they must satisfy Nelson's 's identity [47].
u
V log क, for [0, 1].
(14)
Based on this, we choose the structural forms of u and e at (1z (r. as
a² w(r Te) =
V log (Tr) (1- -1- B.) Dos Dy)),
(15)
(rr, Te) =
V log #. (ra) + (1 (a0--,,)NN((,, De, Dp))
2
where NN(.) is neural network that takes as input the state rt. the step index t. and the represen- tations of the context De U D2 computed from a pre-trained encoder. If we were to directly use the score function of Te, NN could simply output the Gaussian score. However, we instead wrap the inputs with a neural net for added flexibility. Note that NN encompasses this score function as a special case, ensuring compatibility with Nelson's 's identity. Based the choice of the drifts, we set the forward and backward transitions for SMCS as Euler-disssetizaaiion of the SDEs at (7, (
Fo Te-1) = [(r Tc-1 + 1e-1(1;1) re) (r.-1 Fr + hr 2h.1)
2h+1).
(16)
While it is in principle possible to model the entire drifts using neural network, recent empirical studies [24] indicate that approaches without explicit score functions often fail to work for even simple target distributions. By combining score-based gradients with neural proposals, our method stabilizes particle transport and enables efficient and expressive test-time inference. Refer to Fig. 2 to see the overall model structure and procedure of TTSNP. And, see Appendix for details on the neural network architecture and modeling choices.
3.3 Training Objective
To ensure effective sampling with the SMCS framework, we must carefully design both the training objective and the data used train F, and B,. If properly designed, the importance weights may become overly concentrated on a few samples, leading to a severe particle degeneracy, which can negatively impact the diversity and effectiveness of the sampling process. To promote well- balanced importance weights and obtain high-quality posterior samples, along with the resampling, we employ two complementary loss terms: (1) the Kullback-Leibe (KL) divergence loss between two path measures, denoted as LKL, and (2) the log- likeliidddood maximization loss, denoted as LLL.
KL. Divergence Loss. The KI. divergence loss, Lputh. is formulated using two path measures, TP" and P which correspond to the forward and backward SDES, respectively. Intuitively, the forward path measure 7d can be interpreted as the joint distribution Q(rā:r) in the limit as T -> oc [5]. If this divergence is minimized to zero, it implies that the forward and backward transitions perfectly -ine-reversible, thereby ensuring ideal sampling [7]. Consequently, training the transition kernels with this objective directly enhances the accuracy and quality of samples generated by the SMCS procedure. Then, given the sequential nature of our approach, the KL divergence is evaluated independently on each subinterval [Tc-1. Tel. yielding:
T
LKL ED,
Σ
DEL
In
(17)
where KL divergence is averaged by pseudo context due to the randomness D,. This formulation can be rewritten as simplified objective using importance weights, following Chen et al. [7]. For detailed definitions and the full loss derivation, refer to Appendix B.

---

1.5
0.7
0.7
-0.1
-0.
-05
2.5
-17
-2.0
-1.2
-0.4
0.4
-2.5
-2.0
-12
-04
-25
-2.
-13 2.0
Conce pant
1.5
-.1-
93
-55
1.0
Figure Visualization of posterior samples. (left) Inference results using latent samples drawn from variational posterior constructed solely from pseudo representations without access to the true context points; (middle) inference results using latent samples refined via TTSNP's SMCS procedure; and (right) inference results from the Fine-tune baseline.
Marginal Log-likelihood Loss. In addition to the KL divergence loss, we introduce a log- likelihood maximization loss, LLL, which encourages the generated samples to capture diverse modes of the target distribution. This objective designed to enhance the representation of epis- temic uncertainty while reducing the number of samples required for reliable estimation. The loss is defined as Lu E(z,yeE log Σ (ylæ, r7, Dz). where r/ denotes samples obtained from final step of the SMCS procedure. Thus our final training objective becomes = LxI. + LLI.-
Experiments
In this section, conduct series of experiments empirically validate the effectiveness of TTSNPs across a variety of settings, with a particular focus on regression tasks. We utilize two representative models from the latent NP family: simple NP [19]. which is the earliest model in this line of work, and DANP [33], a recent model that has demonstrated strong performance and broad applicability across various tasks. Since our TTSNP is designed be compatible with exist- ing model architectures, we consider the following three baseline approaches: 1) Pre-train: a model that has been pre-trained on the training dataset is directly used for inference the test dataset in zero-shot manner, relying solely on the pre- rre-trained latent variational posterior without any additional training: 2) Fine-ture:
latent path further adapted for inference by training the transition kernel using the available data at test time; 3) SMCS: the samples from variational posterior transported by SMCS with the ULA transition kernels which does not require training. Unless otherwise specified, we fix the number of latent variable samples to 50 across both our method and all baselines to en- sure fair comparison. In every result table, the best-performing metric highlighted all results are averaged over five different random seeds represented with -siima error bars. In the experiment tables, gray box denotes that the evaluation dataset is the same as the training dataset (seen), while yellow box indicates different evaluation dataset (unseen) for the TTSNP.
bold, and
4.1 Simple NP
Sample Quality. We first conducted an experiment to demonstrate that our method, TTSNP, fectively and efficiently transforms samples drawn from the variational posterior to be better aligned with the true posterior. To this end, we pre-trained an NP model on one- dinenniioonnl GP data with RBF kernels. Using additional data from the same kernel, we trained both the 'Fine-tune' baseline and our TTSNP method. We then performed inference on an evaluation set and compared the quality of posterior samples produced by each method. Table 1 clearly demonstrates that TTSNP generates higher-quality latent variables compared to other baselines, even with the same number of samples, resulting in improved likeliiossoodds for both context and target points. This result shows that our

---

Table 1: Log- logl-iiklihood results for Sample Quality (left) and Matern to RQ (right) scenarios.
Model
RBF
Model
Matern
RQ
context
furget
context
target
Pre-train
0.791 1.0.003
0.352
=0.004
Pre-train
0.654 1:0,008
0.150 1.00110
1,045 10.005
0.711
20.065
Fine-tune
0.833 +0.005
0.382
-0.010
Fine-tune
L6600 1:0000
0.178 10.008
1.094
20.001
0.718
20.002
SMCS
SMCS
0.757 1:0.004
0.186 +00006
1.122 400005
0.735
40.002
TTSNP (ours)
0.893 110:001
0.430
00000
TTSNP (ours)
0.753 100004
0.219 300000
1,122 10.001
0.755 200006
learned intermediate distributions and transition kernels generalize effectively not only additional training data but also to the evaluation data set.
Analysis on generated pseudo context Furthermore, Fig. 3 illustrates how TTSNP benefits from the generated pseudo representations to obtain improved latent poste- rior samples. Specifically, the left, middle, and right sub- figures in Fig. 3 respectively show: (left) inference results using latent samples drawn from variational posterior constructed solely from pseudo representations without access to true context points, (middle) inference using la- tent samples refined via TTSNP's SMCS procedure, and (right) inference results from the Fine-tune baseline. As observed in the left figure, latent samples drawn from the variational posterior built on pseudo representations produce predictions and uncertainty estimates that differ from those of the original pre-trained variational poste-
0.43
0.42
0.41
0.40
0.37
0.50
SONES
TEENE touns
Number of Timse Steps
Figure 4: Marginal log likelihood com- parison between SMCS and TTSNP us- ing NP model. The x-axis indicates the number of time steps.
rior. This alternative perspective, introduced by the pseudo representations, becomes part of the intermediate distributions used in TTSNP. As result, during the transition toward the true pos- terior, such diverse possibilities are continually iwegrated-allovinn TTSNP to capture a broader range of posterior modes more sample-efficienty than standard SMCS.
Analysis on trade off between inference cost and the performance We conducted experiments to analyze the trade-off between inference cost and performance at test time. Specifically, we com- pared performance across different numbers of SMC steps, rather than using the default setting of T 10 from previous experiments. As shown in Fig. 4, both SMCS and TTSNP exhibit improved performance as the number of inference steps increases. Notably, TTSNP achieves faster perfor- mance gains compared to SMCS, demonstrating its greater sample efficiency during inference. This behavior is similar to the trend observed in LLMs, where increasing inference cost-such such through more sampling or decoding steps often leads to improved model performance.
Matern to RQ In this scenario, we assume pre-trained NP model that has been trained on GP data generated using the RBF kernel. We aim to perform inference on GP data generated with either a Matern or Rational Quadratic (RQ) kernel. Here, we assume that a small amount of additional training data is available for the Matern kernel, whereas no extra training data is provided for the RQ kernel. It is important to note that both the Matern and RQ kernels can be seen as different generalizations of the RBF kernel. Our expectation is that, if the SMCS framework is well trained to transform latent samples from the variational posterior into the true posterior using the Matern kernel data, the improved inference quality will also transfer to the RQ kernel, which is another gen- eralization of the RBF kernel, despite the absence of direct training data. Table shows the expected results, demonstrating that when the training and validation datasets share underlying characteristics, our learned intermediate distributions and transition kernels enable the NP to generalize effectively to the validation dataset.
4.2 Varying dimensional tasks with DANP
In this section, we design our experimental setup as follows. Leveraging the ability of the DANP model to handle data with varying input dimensions, we first pre-train the model using data gener- ated from 3D and 4D GP with RBF kernels. Then, we generate training datasets from previously unseen ID and 2D GP data with RBF kernels to train both the fine-tuned baseline and our TTSNP model. After training, we evaluate both models on ID GP data generated in two different settings:
8

---

Table 2: Results of the context and target log- likelihoods for the GP regression task using the DANP models. "nD A" in the first column denotes that the n-- GP dataset was used to train for method A.
ID RBF
Inpu: range shill
Hyperparameter range shift
Model
context
target
context
target
context
target
Pre-triin
1,0488 +00001
+-0.002
+00.00
0,2345
+00
0.8136
4:0.000
-0.1971
+00.015
SMCS
1.2725 1.0.001
0,4836
1:0.004
1,1874
1:0.001
0.4226
1:0.003
1.2051 1:0.001
0.1700 +.0.014
1D Fine-tune
1.1863 +0.005
0.5095
+0.001
1,1270
#0.001
0.4663
+00004
0.8594
+:0
0,1084
+0.001
TTSNP (Ours)
1,2729 #0.001
+0.004
1,2639
+00001
0.6220
1,0890
+:0
0,3615
+0.001
Fine-lune
1.1840 1.0.001
0,4150 1.0.003
1,0859
1:0.001
0.3809
1:0.001
0.9853
1.0.000
0.1485 1.0.00
2D TTSNP (Ours)
1.3099 2:00001
0.6026
#0.003
1,2429
1:00000
0.5913 2:00003
1,0414 1:0.000
0.3075 20,000
Table 3: log-likelihood results for context and target values were obtained for image completion tasks using the EMNIST and Corrupted EMNIST datasets with 3 different corruptions
EMNIST
Comupted
Full
Context
Mean
Model
cuntext
Largel
cuntext
larget
Pre-train
1.334 1: 0.003
0.574 : 0.003
1.241 1: 0.003
0.551 : 0.004
SMCS
1.349 0.005
0.591 0.010
1.351
0.002
0.573 0.008
Fine-tune
1,382 +: 0.004
0.601
0.004
1.381
0.002
0.583 +: 0.005
TISNP (ours)
1.350 = 0.002
0.620 0.003
1.351
0.001
0.613 + 0.004
Figure 5: Visualization of the generated mean and covariance from TTSNP for the EMNIST and CelebA image completion tasks. Each column corresponds (left) the full image, (middle left) the given context, (middle right) the predicted mean, and (right) the predicted standard deviation, respectively.
(1) data within the same distribution as the training set, and (2) data with covariate shifts, such as a shifted input range or modified RBF kernel hyperparameters. This setup allows us assess model performance both on in-distribution inference and under covariate shift conditions. The results in Table 2 confirm that the learned intermediate distributions and transition kernels remain effective even when the input dimension changes, as long as there is an underlying shared structure between the training and evaluation
as being generated from the same RBF kernel, or involving shifts in input range or kernel hyperparameters. This demonstrates the robustness of TTSNP under various covariate shift scenarios.
4.3 Image Completion
In the image completion scenario as well, we adopt setup based on the DANP module. First, leveraging the capability of DANP to handle varying output dimensions, pre-train the DANP model on the CelebA [37] image completion task. After this pre-training phase, we perform few- shot fine-tuning of the pre-trained DANP model image completion tasks constructed from the EMNIST [9] dataset. Then we evaluate TTSNP with the EMNIST validation set and the Corrupted EMNIST dataset. Here, used three different image corruption methods and report the averaged results from these corruptions: 1) 'snow', 2) 'flip', and 3) 'brightness'. Table clearly shows that, compared other baselines, TTSNP more effectively adapts to new tasks with changing output dimensions, resulting in improved probabilistic inference. This shows that TTSNP is effective not only in scenarios where the input dimension varies but also in cases where the output dimension changes. As long as the task used to train the NP model shares some underlying features with the target task, TTSNP can leverage a small amount of additional data to successfully enhance inference performance. Also, similar to previous experiments, the learned SMCS procedure generalizes well to the unseen shifted evaluation tasks. Fig. 5 shows the predictive mean and standard deviation of EMNIST and CelebA data from TTSNP.
4.4 Ablation
Ablation on the number of pseudo context points. We conducted an ablation study to investigate the effect the number of pseudo context points on model performance using the NP model for GP regression with an RBF kernel. Our default configuration employs 12 pseudo representations. To analyze the sensitivity to this design choice, we additionally evaluated models with and 48 pseudo representations. As summarized in Table 4 (left), increasing the number of pseudo representations

---

Table Ablation studies on (left) the number of pseudo context points and (right) the recalibrated
variance in constructing intermediate distributions.
Number Representations
(default)
48
Type of Varinnce
Learnt
Fixed
Context
0.800=0.001
0.893+000001
0.895=0.001
Context Taget
1.22=000011 0.46320.003
1.273=0.00
Target
0.421+0.002
0.430++0.001
0.438+0.00)
0.661
consistently improves both context and target performance, suggesting that richer pseudo context provides better guidance for latent inference.
Ablation on the recalibrated variance. We further performed ablation the role of recali- brated variance in constructing intermediate distributions for test-time refinement. The recalibrated variance mitigates the underestimation of uncertainty in the original variational posterior, enabling more accurate and stable inference without additional training. To validate this effect, we compared TTSNP using DANP models on a ID GP regression task using either the learnt variance or recali- brated fixed variance. The results, reported in Table 4 (right), demonstrate that recalibrated variance yields improved target prediction performance, highlighting its importance in our refinement strat- egy.
Ablation on the number of training data Here, we evaluate the learning efficiency and convergence speed of TTSNP compared the Fine-tune base- line by analyzing their target logllikelhoodd perfor- mance. Both methods start from the same pre- trained NP model trained on 1D) GP data generated with an RBF kernel. We vary the number training data batches from 10 to 1000 and measure perfor- mance using the mean and standard deviation the target log-li logllikel across five random seeds. As shown in Fig. 6, TTSNP achieves significantly faster convergence and consistently outperforms the Fine- tune method, even with substantially fewer training samples. This result shows that method not only samples more improved latent variables but is also more efficient train compared the Fine-tune. Refer to Appendix D for additional experiments, including ablation studies on the number of sam- ples and training objectives.
0.4
13lelthood
0.3
0.-
0.1
TTSNP sourg
Fime-lune
10
50
100
500
1000
Number of Training Batch
Figure 6: Comparison of target log- likelihood between TTSNP and the baseline 'Fine-tune' under varying amounts of train- ing data, using the same number of samples.
5 Conclusion
In this paper, we first observed that the latent variational posteriors of existing NP models are of- ten miscalibrated. To address this issue, we proposed TTSNP, test-time scaling method based on a learned SMCS framework that transforms samples from the variational posterior into ones that better match the true posterior. TTSNP achieves this by learning both the construction of inter- mediate distributions through pseudo context representations conditioned on the input context and the approximation of score functions required for transition kernels. This enables TTSNP to guide samples more efficiently and effectively toward the true posterior by leveraging diverse and infor- mative intermediate distributions. Our experiments demonstrate that TTSNP significantly improves posterior inference quality across wide range of settings, including changes in input and output dimensions, covariate shifts, and cross-task generalization scenarios.
Nonetheless, TTSNP comes with some limitations. Training the model requires additional data and introduces extra memory and computational costs. Moreover, unlike standard pre-trained NPs, TTSNP requires additional computation at inference time, such as gradient evaluations for the SMCS procedure. Although this overhead can be mitigated by reducing the number of SMCS steps-allowing for trade-off between efficiency and sample quality-future work may focus on minimizing this computational burden while maintaining strong performance.
10

---

Acknowledgement
This work was supported by Institute of Information & Communications Technology Planning & Evaluation(IITP) grant funded by the Korea goveM NNo.RS-2019-11190000 Artificial Intelligence Graduate School Program(KAIST), Institute of Information & Communications Tech- nology Planning & Evaluation(IITTP) grant funded by the Korea governmen((MSIT) (No.RS-2022. 11220713, Meta-learning Applicable to Real-world Problems), Institute of Information & Commu- nications Technology Planning & Evaluation(IITP) grant funded by the Korea governnent(MSrT) (No.RS-0055-2211931, AT Star Fellowship(Koote Kookmin University)). Artificial intelligence industrial convergence cluster development project funded by the Ministry of Science and ICT(MSIT, Ko- rea)&Gwangju Metropolitan City, and the MSIT(Ministry of Science, ICT), Korea, under the Global Research Support Program in the Digital Field program(IITP- PTTT-2204-RSS- R---222--00417777700 by the IIIs(tnt for Information & Communications Technology Planning & Evaluation).
supervised
This work was also supported by the Institute of Information & Communications Technology Plan- ning Evaluation (IIT)) with a grant funded by the Ministry of Science and ICT (MSIT) of the Republic of Korea in connection with the Global AI Frontier Lab International Collaborative Re- search, and by the Samsung Advanced Institute of Technology (under the project Next Generation Deep Learning: From Pattern Recognition to AI).
This work was also supported by the NIH/NHLBI Award ROIHLI48248, NSF Award 1922658 (NRT-HDR: FUTURE Foundations. Translation, and Responsibility for Data Science), NSF CA- REER Award 2145542, ONR N00014-23-1-2634, NSF 2404476, and Apple.
References
[1] Cesar Almecija, Apoorva Sharma, and Navid Azizan. Uncertainty-aware meta-learning for multimodal task distributions. arXiv preprint arXiv:2210:0188) 2022.
[2] Han Altae-Tran, Bharath Ramsundar, Aneesh S Pappu, and Vijay Pande. Low data drug dis- covery with one-shot learning. ACS central science, 3(4::283-293, 2017.
[3] Jason Ansel, Edward Yang. Horace He, Natalia Gimelshein, Animesh Jain, Michael Voznesen- sky, Bin Bao, Peter Bell, David Berard, Exgeni Burovski, Geeta Chauhan, Anjali Chourdia, Will Constable, Alban Desmaison, Zachary DeVito, Elias Ellison, Will Feng, Jiong Gong, Michael Gschwind, Brian Hirsh, Sherlock Huang. Kshiteej Kalambarkar, Laurent Kirsch, Michael Lazos, Mario Lezcano, Yanbo Liang, Jason Liang, Yinghai Lu, CK Luk, Bert Maher, Yunjie Pan. Christian Puhrsch, Matthias Reso, Mark Saroufim, Marcos Yukio Siraichi. Helen Suk, Michael Suo, Phil Tillet, Eikan Wang, Xiaodong Wang, William Wen, Shunting Zhang, Xu Zhao, Keren Zhou, Richard Zou, Ajit Mathews, Gregory Chanan, Peng Wu, and Soumith Chintala. Py PyTccc 2: Faster Machine Learning Through Dynamic Python Bytecode Trans- formation and Graph Compilation. In 29th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2 (ASPLOS '24). ACM, April 2024. doi: 011444660666664406666 URL httph:ttt//co http://yyttttcccccccooooo org.eehh.rrreeeee assets/ pytorch2- 2.pf.
[4] Matthew Ashman, Cristiana Diaconu, Junhyuck Kim, Lakee Sivaraya, Stratis Markou, James Requeima, Wessel Bruinsma, and Richard Turner. Translation equivariant transformer neural processes. arXiv preprint arXiv:2406. 00022066.2440 /2409, 2024.
[5] Paolo Baldi. Stochastic calculus. Springer, 2017.
[6] Wessel P Bruinsma, Stratis Markou, James Requiema, Andrew YK Foong, Tom R Andersson, Anna Vaughan, Anthony Buonomo, J Scott Hosking, and Richard E Turner. Autoregressive conditional neural processes. arXiv preprint arXiv:2303.14460 2023.
[7] Junhua Chen, Lorenz Richter, Julius Berner, Denis Blessing, Gerhard Neumann, and Anima Anandkumar. Sequential controlled langevin diffusions. arXiv preprint arXiv:2412.07081, 2024.
[8] NICOLAS CHOPIN. Central limit theorem for sequential monte carlo methods and its appli- cation to bayesian inference. The Annals of Statistics, 3286)22385-241, 2004.
11

---

[9] Gregory Cohen, Saeed Afshar, Jonathan Tapson, and Andre Van Schaik. Emnist: Extend- ing mnist to handwritten letters. In 2017 international joint conference on neural networks (IJCNN) pages 2921-2926. IEER, 2017.
[10] C. Dai, J. Heng, P. E. Jacob, and N. Whiteley. An invitation to sequential Monte Carlo sam- plers. Journal of the American Statistical Association, 1 117:1587-1600, 2022.
[11] Pierre Del Moral, Arnaud Doucet, and Ajay Jasra. Sequential Monte Carlo samplers. Journal of the Royal Statistical Society Series B: Statistical Methodology, (8(3):411-4366 2006.
[12] A. Doucet, W. Grathwohl, A. G. D. G. Matthews, and H. Strathmann. Score-based diffusion meets annealed importance sampling. In Advances in Neural Information Processing Systems 35 (NeurIPS 2022), 2022.
[13] Leo Feng, Hossein Hajimirsadeghi, Yoshua Bengio, and Mohamed Osama Ahmed. Latent bottlenecked attentive neural processes. arXiv preprint arXiv:2211.08458 2022.
[14] Leo Feng, Frederick Tung. Hossein Hajimirsadeghi. Yoshua Bengio, and Mohamed Osama Ahmed. Memory efficient neural processes via constant memory attention block. arXiv preprint arXiv:2305. 14567, 2023.
[15] Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adap- tation of deep networks. In International conference on machine learning, pages 1 1135. PMLR, 2017.
[16] Edwin Fong, Chris Holmes, and Stephen Walker. Martingale posterior distributions. arXiv preprint 103.::20...1667,, 2021.
[17] A. Y. K. Foong, W. P. Bruinsma, J. Gordon, Y. Dubois, J. Requeima, and R. E. Turner. Meta- learning stationary stochastic process prediction with convolutional neural processes. In Con- ference on Neural Information Processing Systems, 2020.
[18) M. Garnelo, D. Rosenbaum, C. Maddison, T. Ramalho, D. Saxton, M. Shanahan, Y. W. Teh, D. J. Rezende, and S. M. A. Eslami. Conditional neural processes. In International Conference on Machine Learning, 2018.
[19] M. Garnelo, J. Schwarz, Rosenbaum, F. Viola, J. Rezende, S. M. A. Eslami, and W. Teh. Neural processes. ICML. Workshop on Theoretical Foundations and Applications of Deep Generative Models, 2018.
[20] T. Geffner and J. Domke. Langevin diffusion variational inference. In Proceedings of The 26th International Conference Artificial Intelligence and Statistics (AISTATS 2023), 2023.
[21] J. Gordon, W. P. Bruinsma, A. Y. K. Foong, J. Requeima, Y. Dubois, and R. E. Turner. Convolu- tional conditional neural processes. In International Conference on Learning Representations, 2020.
[22] Ido Greenberg, Shie Mannor, Gal Chechik, and Eli Meirom. Train hard, fight easy: Robust meta reinforcement learning. Advances in Neural Information Processing Systems, 36:68276- 68299, 2023.
[23] Taicheng Guo, Xiuying Chen, Yaqi Wang, Ruidi Chang, Shichao Pei, Nitesh V Chawla, Olaf Wiest, and Xiangliang Zhang. Large language model based multi-agents: A survey of progress and challenges. arXiv preprint arXiv:2402.01688. 2024.
[24] Jiajun He, Yuanqi Du Francisco Vargas, Dinghuai Zhang, Shreyas Padhy, RuiKang OuYang, Carla Gomes, and José Miguel Hernández-Lobato. No trick, no treat: Pursuits and challenges towards simulation-fia free training of neural samplers. arXiv preprint arXiv:2502.06683 2025.
[25] Xin He, Kaiyong Zhao, and Xiaowen Chu. Automl: A survey of the ttate-of-the-art.
Knowledge-bass systems, 212: 2021.
[26] Heng, A. N. Bishop, Deligiannidis, and A. Doucet. Controlled sequential Monte Carlo. The Annals of Statistics, 48(5):2904-2929, 2020.
12

---

[27] Timothy Hospedales, Antreas Antoniou, Paul Micaelli, and Amos Storkey, Meta-learning in neural networks: survey. IEEE transactions on pattern analysis and machine intelligence, 44(9):519955199 2021.
[28] H. Kim, A. Mnih, J. Schwarz, M. Garnelo, S. M. A. Eslami, D. Rosenbaum, and V. Oriol. Attentive neural processes. International Conference on Learning Representations, 2018.
[29] Sunwoo Kim, Minkyu Kim, and Dongmin Park. Test-time alignment of diffusion models with- out reward over-optimization, 2025. URL
org/ abs//aa///aa 2501. 05803.
[30] Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In Interna- tional Conference on Learning Representations, 2015.
[31] Tuan Anh Le, Hyunjik Kim, Marta Garnelo, Dan Rosenbaum, Jonathan Schwarz, and Yee Whye Tch. Empirical evaluation of neural process objectives. In NeurIPS workshop on Bayesian Deep Learning, page 71, 2018.
[32] Hyungi Lee, Eunggu Yun, Giung Nam, Edwin Fong, and Juho Lee. Martingale posterior neural processes, In International Conference on Learning Representations, 2023.
[33] Hyungi Lee, Chaeyun Jang, Dongbok Lee, and Juho Lee. Dimension agnostic neural pro- cesses. arXiv preprini arXiv:2502.20661 2025.
[34] J. Lee, Y. Lee, J. Kim, E. Yang. S. J. Hwang, and Y. W. Teh. Bootstrapping neural processes. In Conference on Neural Information Processing Systems, 2020.
[35] Juho Lec, Yoonho Lee, Jungtack Kim, Adam Kosiorek, Scungjin Choi, and Yee Whye Tch. Set transformer: A framework for attention-based permutation-invariant neural networks. In International conference on machine learning, pages 3744-3753. PMLR, 2019.
[36] Xinzhu Liang, Joseph M Lukens, Sanjaya Lohani, Brian Kirby, Thomas A Searles, and Kody JH Law. Sme is all you need: Parallel strong scaling, arXiv preprint arXiv:2402.06/73, 2024.
[37] Ziwei Liu, Ping Luo, Xiaogang Wang, and Xiaoou Tang. Deep learning face attributes in the wild. In Proceedings of International Conference on Computer Vision (ICCV), December 2015.
[38] C. Louizos. X. Shi, K. Schutte, and M. Welling. The functional neural process. Conference on Neural Information Processing Systems, 2019.
[39] Stratis Markou, James Requeima, Wessel P Bruinsma, Anna Vaughan, and Richard E Turner. Practical conditional neural processes via tractable dependent predictions. arXiv preprint 27X1v:2203.08775. 2022.
[40] Alessandro Mastrototaro and Jimmy Olsson. Online variational sequential monte carlo. In Forty-first International Conference on Muchine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024. URL https:/ ptttp:://openreviee net forumm ddd CpP63CW
[41] Alex Matthews, Michael Arbel, Danilo Jimenez Rezende, and Arnaud Doucet. Continual re- peated annealed flow transport monte carlo. In International Conference Machine Learning, 2022.
[42] Cory McCartan and Kosuke Imai. Sequential monte carlo for sampling balanced and compact redistricting plans. The Annuls of Applied Statistics, 17(4):3300-3323, 2023.
[43] Marko Mlikota and Frank Schorfheide. Sequential monte carlo with model tempering. Studies Nonlinear Dynamics and Econometrics, 28(2):249-269, May 2023. ISSN 1558-3708. doi: 1//nee222222222200000 URL http: /xxxxxxx////////////// doi wow.ooo//........o 1515/ 2015//nma ande- 0022-01033
13

---

[44] Christian A. Naesseth, Scott W. Linderman, Rajesh Ranganath, and David M. Blei. Variational sequential monte carlo. In Amos J. Storkey and Fernando Pérez-Crie editors, International Conference on Artificial Intelligence and Statistics, AISTATS 2018, 9-11 April 2018, Playa Blanca, Lanzarote, Canary Islands, Spain, volume 84 of Proceedings of Machine Learning Re- search, pages 968-977. PMLR, 2018. URL http: http::::/ppoceeddngeee ngs. ml ptrss/rss///4 beesseth1@a.htne
[45]
Radford M Neal. Annealed importance sampling. Statistics and computing, 11:125-139, 2001.
[46] Radford M Neal and Radford M Neal. Monte carlo implementation. Bayesian learning for neural networks, pages 55-98 1996.
[47] Edward Nelson. Dynamical theories of Broonian motion. Princeton University Press, 1967.
[48] Tung Nguyen and Aditya Grover. Transformer neural processes: Uncertainty-aware meta learning via sequence modeling. In International Conference Machine Learning, 2022.
[49] Deep Shankar Pandey. Uncertainty-Aware Meta-Learning for Learning from Limited Data. PhD thesis, Rochester Institute of Technology, 2025.
[50] Jürgen Schmidhuber. Evolutionary principles in self-referential learning, or on leurning how learn: the meta-meta-.. hook. PhD thesis, Technische Universität München, 1987.
[51] A. Thin, N. Kotelevskii, A. Doucet, A. Durmus, E. Moulines, and M. Panov. Monte Carlo variational auto-encoders. In Proceedings of the 38th International Conference on Machine Learning (ICML 2021). 2021.
[52] Surya T Tokdar and Robert E Kass. Importance sampling: review. Wiley Interdisciplinary Reviews: Computational Statisties, 2(1):54 60, 2010.
[53] Francisco Vargas, Shreyas Padhy, Denis Blessing, and Nikolas Nusken. Transport meets varia- tional inference: Controlled Monte Carlo diffusions. In International Conference on Learning Representations, 2024.
[54] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Conference on Neural Information Processing Systems, 2017.
[55] Qi Wang and Herke Van Hoof. Learning expressive meta-representattions with mixture of expert neural processes. Advances in neural information processing systems, 35:26242-26255, 2022.
[56] Qi Wang, Marco Federici, and Herke van Hoof. Bridge the inference gaps of neural processes via expectation maximization. In The Eleventh International Conference on Learning Repre- sentations, 2023.
[57] Xuezhi Wang. Jason Wei, Dale Schuurmans, Quoe Le, Ed Chi, Sharan Narang, Aakanksha Chowdhery, and Denny Zhou. Seff-consistency improves chain of thought reasoning in lan- guage models. arXiv preprint 2203 2022.
[58] Jason Wei, Xuezhi Wang. Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoe V Le, Denny Zhou, et al. Chain-of -hain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824-2498 2022.
[59] H. Wu, J. Köhler, and F, Noé. Stochastic normalizing flows. In Advances in Neural Information Processing Systems 33 (NeurIPS 2020), 2020.
(60] Qiyuan Zhang, Fuyuan Lyu, Zexu Sun, Lei Wang, Weixu Zhang, Wenyue Hua, Haolun Wu, Zhihan Guo, Yufei Wang. Niklas Muennighoff, et al. A survey test-time scaling in large language models: What, how, where, and how well? arXiv preprint arXiv:2503.24235, 2025.
[61] Stephen Zhao, Rob Brekelmans, Alireza Makhzani, and Roger Baker Grosse. Probabilistic inference in language models via twisted sequential monte carlo. In Forty-first International Conference on Machine Learning, ICML. 2024, Vienna, Austria, July 21-27, 2024. OpenRe- view.net, 2024. URL https: httpp::::////ppeevviieee ew. net/ forum? d= f A AONNES1
14

---

[62] Zheng Zhao, Sebastian Mair, Thomas B Schön, and Jens Sjölund. On feynman -ac training of partial bayesian neural networks. In International Conference on Artificial Intelligncee and Statistics, pages 3223-3231. PMLR, 2024.
15

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our main contribution is to fix the samples from the uncalibrated latent varia- tional posterior using a learned SMC framework. And we propose a method which clearly aligned with this objective.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification:
In the 5, we discuss about the limitation TTSNP.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The au- thors should reflect on how these assumptions might be violated practice and what the implications would be.
The authors should reflect the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the ap- proach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a areech- speech-to-text system might not be used reliably provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to ad- dress problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
16

---

Answer:
[NA]
Justification:
We don't t have any theoretical results for this paper.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theo- rems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch to provide intuition.
Inversely, any informal proof provided in the of the paper should be comple- mented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We fully disclose all necessary information for reproducing the main experi- mental results, including dataset descriptions, experimental setup, evaluation metrics, and hyperparameters, clearly presented in Appendix C.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided not.
If the contribution is a dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, may be necessary to either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how replicate the results, access a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate the research performed.
While NeurlPS does not require releasing code, the conference does require all sub- missions to provide some reasonable avenue for reproducibility, which may depend on the nature the contribution. For example
(a) the contribution primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results way produce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case au- thors are welcome to describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
17

---

5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We provide detailed experimental setting in Appendix C, and provide the original codebase within the supplementary matrial.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https: //npp.s.co
publ c guides CodeSulomi s conPolicy cy) for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (ht (https: //niss-. cc/ publll c/ /ggu/uuia quudde
des (Codeeubbi as i sionooo ol comPolicy icy) for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: Yes, we provide such details in Appendix C.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropri- ate information about the statistical significance of the experiments?
Answer:
[Yes]
Justification: All the experiments listed averaged over five different random seeds, and we provide standard deviation as well.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
18

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report I-sigma error bars, but one should state it. The authors should prefer- ably report 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
Computational resource details are provided Appendix
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics https: /nwuriiusspp.. cc cc/publice wtw.cctiiiiieeeeeeeeeeeee line??
Answer:
[Yes]
Justification:
Our research fully complies with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No. they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (c.g., there is a special consid- eration due to laws regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Ye]]
Justification: Yes, we provide full broader impacts in Appendix E.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
19

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact spe- cific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitiga- tion strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [Yes]
Justification: We describe appropriate safeguards for responsible use and release of research, including controlled access and clear usage guidelines, detailed in Appendix E.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have a high risk for misuse or dual-use should be released with necessary safeguards allow for controlled use of the model, for example by quiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification:
All datasets and models utilized are clearly listed in 4 and Appendix C,
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
20

---

If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, papersw thoode com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is available online, the authors encouraged to reach out the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documenta- tion provided alongside the assets?
Answer: [Yes]
Justification: Assets of our TTSNP framework (including code) are clearly provided in 4, Appendix C, and supplementary materials.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can cither create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsoureing experiments and research with human subjects, does the pa- per include the full text instructions given to participants and screenshots, applicable, well as details about compensation (if any)?
Answer: [NA]
Justification: Our study does not include crowdsourcing or human-subject experiments.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contri- bution of the paper involves human subjects, then much detail as possible should be included the main paper.
According to the NeurlPS Code of Ethics, workers involved in data collection, cura- tion, other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed the subjects. and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
We do not have any issues regarding IRB approvals.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects
Depending on the country in which research is conducted, IRB approval (or equiva- lent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
21

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology. scientific rigorousness, or originality of the research, declaration is not required.
Question:
Answer: [NA]
Justification:
We used LLM only for grammar checking purposes.
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https: / /neur i ps cc/Conferrnccss///////oooesocs//
2025 (025/LLM) for what should or should not be deseribed.
22

---

A Related Works
Neural Processes The original Conditional Neural Process (CNP) [18] introduced simple encoder-decoder architecture based on multilayer perceptrons. Its extension, the Neural Process (NP) [19], incorporated a global latent variable to model epistemic uncertainty. In this paper, they first suggest the ELBO objective to train the NP models with a latent variable. This design was further improved by Attentive Neural Processes (ANP) [28], which integrated attention mechanisms into the encoder for more expressive context representations. On the other hand, Gordon et [21] and Foong et al. [17] incorporate convolutional deep set module to introduce translation equívari- ance into NP modules. However, since these models rely on convolutional computations, they are structurally limited and difficult to extend to data beyond three dimensions. Recently, Transformer Neural Processes (TNP) [48] replaced MLPs or attention layers with masked transformer layers to enhance scalability and long-range interaction modeling. Building TNP, the Dimension-Agnosee Neural Process (DANP) [33] introduced the Dimension Aggregator Block (DAB), enabling a single model to handle tasks with varying input-output dimensions. Our proposed test-time scaling via SMC is developed on top of latent- tatent-path-Bee path-based models such as NP and DANP.
Beyond these, many works have explored improved uncertainty quantification, expressivenes, and inference strategies. For instance, Functional Neural Processes (FNP) [38] used local latent vari- ables to better capture function-specific variations. Bootstrapping Neural Processes (BNP) [34] introduced a residual bootstrapping mechanism to address model misspecification, while Martin- gale Posterior Neural Processes (MPNP) [32] proposed martingale posterior distribution [16] the- ory based Bayesian inference for more principled uncertainty estimation. Recent variants like Transiation-Equiivriaat Transformer Neural Processes (TE-TNP) [4], Latent Bottlenecked Atten- tive Neural Processes (LBANP) [13], and Mixture-of-Expert Neural Processes (MoE-NP) [55] im- prove generalization and scalability through architectural innovations such as equivariant design, latent bottlenecks, and mixture modeling. Other contributions include Autoregressive Conditional Neural Processes (AR-CNP) [6] for sequential prediction, Self-normalized Importance-weightee Neural Processes (SI-NP) [56] for refined inference, Constant Memory Attentive Neural Processes (CMANP) [14] for memory-efficient modeling, and Gaussian Neural Processes (GNP) [39] for tractable predictive covariance modeling.
Sequential Monte Carlo Our work draws upon foundational concepts in SMC methods and their integration with learned stochastic dynamics. While SMC is designed to efficiently sample from se- quential distributions, it has also been extensively explored for sampling complicated posterior dis- tributions through sequential tempering [43], introducing learnable variational distributions [44, 40], or incorporating Langevin dynamics to replace sequential transitions [7]. In parallel, recent works have focused on inference-time scaling in foundation models such as LLMS and diffusion models, leveraging SMC to facilitate this scaling. Zhao et al. [61] treat the autoregressive prediction of LLMSs as an SMC process and define conditional target distribution with a learnable twisted function, en- abling the guiding of pretrained LLMs to handle new tasks. In the diffusion models, Kim et al. [29] enhance sample efficiency by using a tempered SMC sampler that balances exploration and exploitation during the sampling process, achieving superior performance in tasks such as aesthetic reward optimization and multi-objective optimization.
B Details about SMCS and Modeling Transition kernels
B.1
Sequential Monte Carlo Samplers
In this section, we briefly review Sequential Monte Carlo Samplers (SMCS) [11], which form the foundational framework for our test-time scaling approach. For a more comprehensive overview, we refer the reader to Del Moral et al. [11], Dai et al. [10].
Let 7 (Ix) := / denote a target distribution, where the unnormalized density >xx) is known pointwise but the normalization constant Z is unknown. SMCS aims to approximate expectations with respect to ग, such as Ex[ʃ] for a test function f. while simultaneously providing an estimate of Z. Rather than sampling directly from the target SMCS constructs sequence of intermediate distributions { beginning with tractable initial distribution To and gradually transporting toward the target TT- π. These intermediate distributions are designed to interpolate between To
23

---

Algorithm Multinomial Resampling
Require: Particles {at] time step and corresponding importance weights {w])
Ensure: Resampled particles (æt) and normalized importance weights {wit j=1--
1: Compute the normalized importance weights: viद
wi for i E NJ
Σi-1 w
2: Compute the Effective Sample Size using wis: ESS (
3: if ESS < 0.3N then
4: for :ori=1 1 N do
5: Sample j from {_...., ] using multinominal distribution with probability {
6: Replace particle =r with o = m!! end for
7:
8: Replace normalized importance weights wis with पj for i .J
9:
else
10:
Keep particles { and normalized importance weights (w)
11:
end if
and т, typically becoming increasingly complex as increases. Since the intermediate distributions generally not analytically tractable, SMCS employs sequential importance sampling adjust the samples accordingly.
In detail, let yt be the unnormalized density of TT, i.e., (x) (Z) for t = . (note that Zo = as To is assumed to be fully known). SMCS introduce forward transition kernels (Fr Ie- xe-1)) which propagates sample _T-1 to zs. Starting with zo no, sample path To:T = (ro....,, xr) is generated from the joint proposal with density,
T
Q(Fo.) == (x0))) II Fe x2-1))
(18)
In principle, one would like to use the marginal f ( ) as proposal density for the target æ, but this generally intractable. Instead, SMCS augments the target to the path space Tu:T with sequence of backward transition kernels (B.-1II 2t) and defines the unnormalized path density
(T(or)) = I( where [( = (xTT) II
T
B B
(19)
t=1
sample path zo.r drawn from Q is then corrected with the importance weights,
wr(#or) = I'(2o.) (xT) 1r-1 xt) Q(tor)) To(2o) II Fe(x; xt-1)
(20)
Given the forward and backward kernels, SMCS starts by drawing N i.idd. particles [= from To. At each iteration, the particles sequentially extended by drawing samples from the forward kernel Fr, and the extended particles are re-weighted through the importance weights, which are computed in recursive fashion as follows:
z1) 7-1 (((c) x)-1)
B1
wo (=())) = 1, (x ut wi-1 1
for 1,..., (21)
After completing the sequential update up to the step T, we have an estimator for the expectation E_ and the normalization constant Z,
E-If ~
1 N
Σ
wr (T((r))) f(oT), Ey= wr(x
Z =
1 N
Σ wr (
(22)
Vanilla SMCS can suffer from particle degeneracy: after several iterations, only handful of par- ticles carry nearly all the weight while the rest contribute little. standard remedy resampling, in which new particle set is drawn from the current importance-weighted empirical distribution.
24

---

In practice, resampling is triggered adaptively, for instance, whenever the effective sample size falls below chosen threshold. When the resampling applied, the recursive importance weight updates, and the estimators for expectations and the normalizing constant require minor adjustments. Refer to Algorithm to see how we conducted resampling this paper. For full treatment, see Del Moral et al. [11].
Choice of intermediate distributions. There are various options for designing intermediate dis- tributions; see Del Moral al. [11]. Dai al. [10] for examples. A common choice is the geometric annealing path setting Yo: (z) = To(=)² (1)))) with coefficients 0 = 30 < B < BT
Choice of transition kernels. The design of forward and backward kernels a critical component of SMCS. A widely used approach is to set F. as transition invariant T/s such as Metropolis- Hastings steps, as proposed Annealed Importance Sampling (AIS) [45]. While there exist optimal backward kernels for given forward kernels [11], they are often intractable in practice. Hence, back- ward kernels are often chosen for tractability. Notably, in AIS, the forward and backward kernels are coupled in a specific way, allowing simplified importance weight computation that does not even require evaluating the forward transition densities. Alternatively, one may consider proposals based on the Unadjusted Langevin Algorithm (ULA), where the forward transition kernels are defined via a discretization of Langevin diffusion [26, 59, 51],
(xx | xx-1) = (x zz-1 + le V log TH (x+-1)), 2h,
(23)
with hú denoting the step size. When he is small, the forward proposal is approximately time- reversible, making the following a natural choice for the backward kernel:
B ( xx))
F x4) N(x,-1 xt + he V log N. (x.), I1,
(24)
One may also consider underdamped Langevin that incorporates auxiliary momentums [20].
Choice of resampling threshold We perform resampling when the effective sample size (ESS) falls below 30% of the total number of particles N, i.c. when ESS < 0.3.V. This threshold is conventional choice widely adopted in SMC-based inference methods [53, 7], as it provides a good balance between computational efficiency and sample diversity. We follow the same criterion in our implementation to maintain stability and prevent particle degeneracy during the sequential inference process.
Learning for SMCS. Since both the intermediate distributions and the forward/backward kernels are crucial design choices in SMCS, it is beneficial to learn them when possible. Learning for SMCS can be formulated as a variational inference problem, where the ELBO provides lower bound on the marginal likelihood defined by the SMCS procedure. A common approach is to begin with un- adjusted Langevin diffusions and modify the drift term (involving the score functions) to maximize the ELBO [12, 20, 53, 7]. Among these, Chen et al. [7] introduced an optimal control framework for tuning continuous-timg extension of SMCS with resampling, which forms the foundation of our proposed algorithm.
B.2 Theoretical Justification of the SMC Refinement
While our approach is primarily motivated by empirical observations, its underlying rationale theoretically supported by established results in the SMC literature. In particular, SMC methods are known to construct a sequence of intermediate distributions that progressively bridge a tractable initial distribution (e.g., a variational posterior) and the true posterior through iterative resampling and transition steps [11, 8J. Under suitable regularity conditions, the weighted particle approxima- tion produced by SMC converges asymptotically to the target posterior as the number of particles increases.
In our setting, initializing SMC with the variational posterior provides reasonable, though often underdispersed, starting point. The sequential refinement process then mitigates this underdispersion by gradually correcting the approximation mismatch through resampling and transition dynamics. Our method is directly grounded in these established convergence guarantees, providing principled justification for using SMC. to enhance posterior quality in latent-variable models.
25

---

B.3 Conceptual Connection to Test-Time Scaling
Test-time scaling (TTS) refers to a family of techniques developed for large language models (LLMs) that enhance inference quality by increasing the computational scale during inference [60]. These approaches include parallel scaling methods such as self-consistency [57] and multi-agent decoding [23], which generate and aggregate multiple reasoning paths, as well as internal scaling strategies like Chain-of-Thought (Co'T) prompting [58], which expand the reasoning depth or struc- ture of single response. Rather than relying on single greedy decoding trajectory, these methods leverage multiple reasoning traces or extended generation paths to either aggregate or select more plausible outputs, ultimately improving inference reliability and accuracy.
Our proposed approach, TTSNP, builds on the same core idea: increasing computational effort at test time to refine inference. Specifically, this is achieved by using larger number of SMC steps or by sampling additional latent variables, both of which serve enhance the quality of posterior approximation. This conceptual parallel motivates the use of the term "scaling" in Test- Time Scaling for Neural Processes, highlighting how our method extends the notion of test-time refinement from LLMs to probabilistic latent- variable models.
B.4 Model structures
As discussed in $ 3.2, our TTSNP framework models both the pseudo representation generator h and the Forward/backward transitions u and v using neural transition kernels. In this section, we provide detailed descriptions of how these two neural modules are implemented.
We begin with the pscudo representation generator h. As noted in $ 3.2, h is designed to be permutation-invarian to ensure exchangcability of the generated context representations. Specif- ically, it takes the context set representations R. and a set of random noise variables c ~ p(c) (sam- pled from a standard Gaussian distribution in our experiments) as input and produces randomized pseudo context representations.
Following Lee et al. [32], we implement h using the induced self-attention block (ISAB) [35], cormutationninvariinn attention module. The pseudo representation Rd generated by ISAB is defined as:
ISAB(e) = MAB(e, where = MAB(Re: c),
where MAB denotes Multihead Attention Blocks [54], and R. the encoded context representation.
Next, we describe the structure of the drift functions u and v, which follow the formulation in Eq. 15:
σ²
u(re, Te) V log के (rr + (1 (,) De,
(25)
2
v t. Te) = 2
VV (o (1 B:)NN(T,, De, DD)
(26)
where NN(.) is a neural network that takes as input the latent state It, the time step t, and the combined representations of the context and pseudo context sets D. U Dpp obtained via the pre- trained encoder. Including t as input to NN helps model the adaptive variance of the intermediate posterior q( D. Dp))
The structure of NN is implemented as a 3-layer Multi-Layer Perceptron (MLP), with the time step t embedded using sinusoidal positional encoding (PE) in Vaswani et al. [54]. The final form of the network is:
NN t, De. Dp) MLP (concat(rr, Re, Rp)) PE(t/T)),
where MLP refers to a 3-layer feedforward network, and PE(.) is the sinusoidal positional embed- ding.
B.5 Details on KL divergence loss
Path measure and Radon-Nikodym Derivative The KL divergence loss, Lpath- is formulated using two path measures, 1a and PP, which correspond to the forward and backward SDEs, respec- tively. Intuitively, the forward path measure pu can interpreted as the joint distribution Q(r).r) in the limit as T [5]. If this divergence is minimized to zero, implies that the forward
26

---

and backward transitions are perfectly time-reversible, thereby ensuring ideal sampling [7]. Conse- quently, training the transition kemells with this objective directly enhances the accuracy and quality of samples generated by the SMCS procedure. A key requirement for enabling principled training and inference the ability to compute the likelihood ratio wo between the forward and reverse- time processes, known as the Radon-Nikodym derivative (RND) [53, 7J. This quantity serves as the backbone for defining meaningful objectives and ensuring the correctness of sample-base ap- proximations. Without tractable form of the RND, learning reliable transition dynamics becomes fundamentally intractable. Fortunately, recent advances [53, 7] provide a computable expression for this erucial quantity as in the following lemma, which leverage directly in our framework.
Lemma B.1 Radon-Nikodyee Derivative [53]). Let Pi [a.0] and P [a.8] denote the puth space measures
of the solutions the forward and backward SDES defined in Eq. 11 and Eq. 12, respectively, over
the time interval [a, b] C [0, 1), where the drift functions u and v satisfy Nelson's 's identity. Assume that the marginal distributions satisfy ra Ta and Sy Tb. Then, Pv (a.1) almost surely, the RND between these two measures can be computed as:
dle [a.b]
(r) = (o In (Ta) {"
( ||u|2 1 V
VoVlog 1og
In (a,bb In dluu [
(re, t)dt
202
f u o²v log (re, dr [- 1 ( dr
As in Lemma 1. we compute by divide trajectories r" and s" into subtrajectories using multiple chunks as follows:
In w In In [ra.T:]
Ipee P Pl [T0mm]]
Pe
Pe mrl In W[rom]] +... + In (
(27)
P
Fr 1.1-)
Following Vargas et al. [53] and Chen et [7], we can further simplify the continuous-time RND expression from Lemma B.1 by applying a discretization approximation. This yields practical discrete-time formulation suitable for implementation:
dIpe
[T0-1.75]
In W/ (..aa....)
(28)
dlPu
[Tn 1.77]
In (T((n-(((()))))) 1)L++ (((nn1))+i---
In Fr, (Fr.) In Tr., (NTw
(29)
In In ()--)) (
where L the number of steps per subtrajectory and is the step size. Here, by the Malryaaa discretization, forward and backward transition densities are computed as follows:
9n--))+ts-11 (Fnew (Fpee) (rrew (rre hhu(Frpee pre, a'), ho (a'))
(30)
P Pin-1)++i (now) ) = (r Fnow + V1(o"V log π u) (I 8'), hb)))
(31)
where a' and b' are ((n - 1)L + i 1)/ and ((n - 1) h)h, respectively.
KL divergence loss Then, given the sequential nature of our approach, the KL divergence is evaluated independently on each subinterval Te). yieldin::
T
LKL. ED
Σ
DKL.
P m
(32)
where KL divergence averaged by pseudo context due to the randomness D,.- This formulation can be rewritten simplified objective using importance weights, following Chen et al. [7]. Following the approaches proposed in Matthews et al. [41] and Chen et al. [7], the KL divergence on each interval can be calculated as:
DxL.
[p
[p
P ||P| ( Epaap
[log W/n (r7((")) (r"))
(33)
27

---

Algorithm Overall TTSNP inference algorithm
Require: Context dataset De., trained pseudo context generator h, and learned forward and back- ward transitions and v.
Ensure: Updated latent particles {ri L.V and corresponding normalized importance weights
First Stage: Generate pseudo Representation
2: Sample noise ~ p(e)
3: Generate pseudo representations Rp using h (Dz,,
4: Second Stage: Generate initial latent particles and importance weights
5: Sample N number of latent particles J from variational latent posterior ((r(Dc)
6: Set corresponding initial importance weights {wi (woJi=1
L.W where w/ N for i {1,-.., N]
7: Third Stage: SMCS procedure
8: for to T do
9: for i to N do
Sample rt WW(ffrf- + huu(ri-1: 1h,1) 11: end for
Update importance weight based on Eq. 8 and Eq. 29 13: Resample the latent variables r and update úi following Algorithm 1
end for
15:
Forth Stage: Compute predictive distribution
With final {r and {w) compute the predictive distribution similar to Eq. 38
where тк T--1 refers to the most recent resampling time prior to Tt-1- Finally, the KL loss CKL can be expressed in practice using importance weights as:
CKL. ED, Σ
detach(w!! log w (
(34)
where detach(-) indicates that gradients are not propagated through the resampled weights w In our experiment, for efficient training, we trained the model using single set of Dp.
B.6 Overall TTSNP latent variable sampling algorithm
In this section, we present the full inference procedure of TTSNP as an algorithm that uses an SMCS procedure refine latent samples toward the true posterior distribution. As outlined in Algorithm 2, the overall method can be divided into four main stages: (1) 'Generate pseudo representation', (2)
'Generate initial latent particles and importance weights', (3) 'SMCS procedure', and (4) Compute predictive distribution'.
In the first stage, 'Generate pscudo representation', we use the pseudo context generator h to create pseudo representations conditioned on the context set. Next, in 'Generate initial latent particles and importance weights', we sample initial latent variables from the variational posterior p(r De) and assign initial weights accordingly. These latent particles are then used to initialize the SMCS procedure.
The third stage, 'SMCS procedure', iteratively updates the latent particles using the learned interme- diate constrrteeee using both the context and pseudo representations- the learned forward and backward transition kernels u and v. Finally, 'Compute predictive distribution' we use the updated latent particles and their associated importance weights to compute the final predic- tive distribution over target outputs.
C Experimental Details
To support reproducibility, we provide our full experimental code part of the supplementary material. Our implementation is based on the official codebase' of DANP [33], and all experi- ments were performed using Py Torch [3]. Training and evaluation were carried out on either an
28

---

Table 5: Model structure details of NP
CATEGORY
DETAILS
MODEL SPECIFICATIONS
HIDDEN DIMENSION FOR DETERMINISTIC PATH
128
HIDDEN DIMENSION FOR LATENT PATH
128
MLP DEPTH FOR DETERMINISTIC PATH
MLP DEPTH FOR LATENT ENCODER PRE-MODULE
MLP DEPTH FOR LATENT ENCODER POST-MODULE
2
DECODER DEPTH
3
NUMBER OF PARAMETERS FOR ID GP REGRESSION 232194
Table 6: Model structure details of NP with TTSNP
CATEGORY
DETAILS
MODEL SPECIFICATIONS
HIDDEN DIMENSION FOR DETERMINISTIC PATH
128
HIDDEN DIMENSION FOR LATENT PATH
128
MLP DEPTH FOR DETERMINISTIC PATH
4
MLP DEPTH FOR LATENT ENCODER PRE-MODULE
4
MLP DEPTH FOR LATENT ENCODER POST-MODULE
2
DECODER DEPTH
3
NUMBER OF MULTIHEAD ATTENTION BLOCKS IN ISAB
2
OUTPUT DIMENSION FOR EACH BLOCK
128
HIDDEN DIMENSION FOR GRADIENT ESTIMATOR
256
MLP LAYER FOR GRADIENT ESTIMATOR
3
NUMBER OF PARAMETERS FOR ID GP REGRESSION
562818
NVIDIA GeForce RTX 3090 or an RTX A6000 GPU. We optimized all models using the Adam optimizer [30], combined with cosine annealing schedule for the learning rate.
Unless stated otherwise, we selected hyperparameters based on validation log- likeliolddood across tasks, using the following search spaces: learning rates from {5 10 5,7 10-5.0 5,9 10 10 3 x 10-4,5 5 x 10- weight decay values of (0, x 10 10-5), and batch sizes of {16, 32).
C.1 Model structural details
In this section, we summarize the architectural details of NP, DANP in Table 5 and Table 7, and their respective variants when integrated with the TTSNP method in Table and Table For the NP and NP+TTSNP models, the description is based on the ID GP regression task. In contrast, both DANP
29

---

Table Model structure details of DANP
CATEGORY
DETAILS
MODEL SPECIFICATIONS
HIDDEN DIMENSION FOR LINEAR PROJECTION IN DIMENSION AGNOSTIC BLOCK
32
HIDDEN DIMENSION FOR SELF-ATTENTION IN DIMENSION AGNOSTIC BLOCK
32
HIDDEN DIMENSION FOR LATENT PATH TRANSFORMER
64
NUMBER OF LAYERS FOR LATENT PATH TRANSFORMER
HIDDEN DIMENSION FOR LATENT PATH SELF-ATTENTION
64
HIDDEN DIMENSION FOR LATENT PATH MLP LAYERS
128
NUMBER OF LAYERS FOR LATENT PATH MLP LAYERS
HIDDEN DIMENSION FOR DETERMINISTIC PATH TRANSFORMER
128
NUMBER LAYERS FOR DETERMINISTIC PATH TRANSFORMER LAYERS
NUMBER OF HEADS FOR DETERMINISTIC TRANSFORMER LAYERS
DECODER DEPTH
2
NUMBER OF PARAMETERS
334562
and DANP+TTSNP are designed maintain a consistent structure and parameter count, regardless of changes in input or output dimensionality.
that TTSNP extends the base latent NP architecture by introducing two additional components: pernutation-invarian pseudo representation generator using the ISAB module [35] and gradient estimator using a Multi-Layer Perceptron (MLP) structure that approx- imates the score function corresponding to the variational posterior constructed from the pseudo representations.
key aspect
C.2 Details on the "Fine-tune"
As outlined in our main experimental setup, we begin with an NP model pre-trained on a large dataset using approximate ELBO objective. Following pre-training, observe that the resulting variational posterior often deviates from the true posterior, particularly few-shot or distribution- shifted settings.
To address this, we introduce a transition kernel and pseudo representation generator trained using small amount of additional data. These components refine the latent samples while keeping the pre-trained NP model frozen, ensuring that the core representation knowledge remains intact and enabling robust adaptation new tasks.
For fair comparison, "Fine-tune" baseline uses same additional data but updates only the latent path of the NP model, with all other components kept frozen. This setup ensures that both methods preserve the pre-trained representation while improving posterior quality through different mechanisms.
30

---

Table 8: Model structure details of DANP with TTSNP
CATEGORY
DETAILS
MODEL SPECIFICATIONS
HIDDEN DIMENSION FOR LINEAR PROJECTION IN DIMENSION AGNOSTIC BLOCK
32
HIDDEN DIMENSION FOR SELF-ATTENTION IN DIMENSION AGNOSTIC BLOCK
32
HIDDEN DIMENSION FOR LATENT PATH TRANSFORMER
64
NUMBER OF LAYERS FOR LATENT PATH TRANSFORMER
HIDDEN DIMENSION FOR LATENT PATH SELF-ATTENTION
64
HIDDEN DIMENSION FOR LATENT PATH MLP LAYERS
128
NUMBER OF LAYERS FOR LATENT PATH MLP LAYERS
HIDDEN DIMENSION FOR DETERMINISTIC PATH TRANSFORMER
128
NUMBER LAYERS FOR DETERMINISTIC PATH TRANSFORMER LAYERS
NUMBER OF HEADS FOR DETERMINISTIC TRANSFORMER LAYERS
DECODER DEPTH
2
NUMBER MULTIHEAD ATTENTION BLOCKS IN ISAB
2
OUTPUT DIMENSION FOR EACH BLOCK
64
HIDDEN DIMENSION FOR GRADIENT ESTIMATOR
192
MLP LAYER FOR GRADIENT ESTIMATOR
3
NUMBER OF PARAMETERS FOR ID GP REGRESSION
475490
C.3 Evaluation Metric for the tasks
Following the approach of Le al. [31], we evaluate baseline models, such NP and DANP, using the normalized predictive bigllikelhood, which estimated through Monte Carlo sampling as:
log logp(1; J. Dzc)
log
Pluj, p(Us,AlZz,111
(35)
kEt
kE
k=1
K
logp(y,A/Ty.. Dj,e)
log
(36)
kEc
kEc
k=1
(k) where r, are sampled from the variational posterior (0,,1D,e) and and denote the context and target points, respectively. However, for the SMC and TTSNP, we utilize our importance weight ws when computing the predictive likelihohood follows:
K
log 1e Dj.0)
log
Σ
(k)
(37)
i
Ret
ket
A=1
log 1og Dj
log
Σ (k)
p(y
(38)
c
c
kge
kee
where i are the normalized importance weight for each sample r (k)
after the SMC procedures.
31

---

C.4 Dataset details of n-dimensional GP Regression task
To construct GP tasks for our experiments, we generate synthetic datasets using GPs equipped with one of three commonly used kernels: the RBF kernel, the Matern 5/2 kernel, and the RQ kernel. Each kernel introduces different inductive biases, allowing us evaluate model robustness across diverse range of smoothness conditions and correlation structures.
The RBF kernel is defined
(-1
[xx - x'133
=2 exp
expp
2P2
while the Matern 5/2 kernel is given by
v5d
5d
x(x,x')= = 1+
where d = 1|x- xxxl.
3
Both use randomly sampled output scales Unif 1.0) and lengthscales E ~ Unif(0.1, 0.6),
and additive Gaussian noise drawn from ~ Unif(0.1, 0.5) is applied to the outputs for numerical stability.
The Rational Quadratic kernel, which generalizes the RBF by integrating over a distribution of lengthscales, is defined
x(x,x)= = (1 1+ xx x
2a02
where we fix the mixture parameter a = 1.0. For its implementation, the lengthscale and output scale sampled similarly from uniform distributions, specifically:
0.1 + (0.6 0.1) Uni Unifoom(0, 1), s s 0.1 + (1.0 - 0.1) Uni form(0, 1).
Given a batch of inputs x, the pairwise normalized distances are computed and scaled accordingly, and small diagonal term o? I added to the covariance matrix to ensure numerical stability.
In our ni-dimensional GP regression setup, we design the data generation process to scale with input dimensionality, ensuring meaningful predictive inference across varying dimensions.
And for the number of context points lcl. we used the sampled number from the range Unif(5n2, 50n - [c]) where indicates the dimension for the GP task. This quadratic scal- ing with n reflects the increased data requirements for higher-dimensional input spaces. Similarly, the number of target points [t] is sampled from Unif (5n2, 50n² - [c)), maintaining fixed upper limit on the total number of points per task. And, inputs x E RR" for both context and target sets are drawn independently from uniform distribution over [-2
This flexible GP data generation process allows us to simulate tasks with varying degrees of smooth- ness, structure, and input dimensionality, providing rigorous testbed for evaluating uncertainty- aware inference methods.
Input range shift As mentioned earlier, we conducted experiments using inputs sampled from the range 2, To evaluate how well TTSNP and baseline methods generalize under covariate shift, we introduced an input range shift scenario. Specifically, modifying only the input domain while keeping all other GP generation settings identical to those used with the RBF kernel. In this shifted setting, inputs x were sampled from [-3, 1]. creating a dis- tributional mismatch between training and evaluation. We then assessed the inference performance of each method this shifted validation set to examine their robustness to covariate shifts in the input space.
constructed validation set by
Hyperparameter range shift In addition to the input range shift, we also evaluated the general- ization ability of TTSNP and baseline methods under a kernel hyperparameter shift, another form of covariate shift where the statistical properties of the data generation process change between training and inference. Specifically, we designed validation scenario in which only the range of hyperpa- rameters used the RBF kernel was altered, while keeping other settings-sun as the number of data points and the input range- identical to the original training setup.
32

---

During training, RBF kernel parameters were sampled from s (01, and ~ Unif(0.1,,.6) 0.6). For the shifted validation set, we extended these ranges to s On1 2.0) and l~ 1n., 1.0), effectively increasing the variance and smoothness of the generated functions. We then performed inference on this new validation data to assess how well each method adapts when faced with unseen kernel configurations, thus testing robustness under structural distribution shifts.
C.5 EMNIST Dataset
For our experiments, we constructed a modified version of the EMNIST Balanced dataset [9]. a widely used benchmark derived from the original NIST Special Database. The full dataset includes 47 alphanumeric character classes, but we curated a 10-class subset for focused evaluation. From this selection, we sampled 24,000 images for training and 4,000 for testing.
Each image is grayscale digit letter rendered in a 28 x 28 resolution, represented as a single- channel input. Pixel locations were linearly scaled to lie within the range [-0.5, 0.5]. and intensity values were normalized to the same interval. During training, for each episode, we randomly chose the number of context points |c| from a uniform distribution over [5, 45). and the number of target points was drawn from Unif(5, 50 - lcl), ensuring total of at most 50 points per task.
Corrupted EMNIST Dataset For the Corrupted EMNIST dataset, we evaluated model robustness by applying three types of image corruptions to the EMNIST evaluation set: (1) snow, flip, and (3) brightesss.
In the snow corruption, we randomly added white noise the image by setting subset of pix- els to the maximum intensity. Specifically, each pixel was independently selected to be corrupted with probability proportional to the corruption intensity. Formally, binary mask was generated where each pixel had a probability 0. obtained by replacing the masked pixels with maximum brightness.
of being set to 1, and the corrupted image was
In the br b ghtness corruption setting, we simulated intensity- -based corruption by amplifying the pixel values and clipping them to stay within valid bounds. Specifically, applied the transforma- tion
data clamp(data x (1 + intensity), 0, 1),
which increases overall brightness proportionally to given intensity factor, followed by clipping to ensure pixel values remain in the 1] range.
For the flip corruption, to simulate geometric distortions, we randomly flipped the image along spatial dimensions. Each image has 5% chance being flipped horizontally and another indepen- dent 50% chance of being flipped vertically. This introduces structural variations while preserving pixel intensity distributions.
C.6 CelebA Dataset
For experiments involving natural image completion, we used the CelebA dataset [37], a large- scale face dataset commonly used in generative modeling benchmarks. The dataset consists of 162,770 training images, 19,867 for validation, and 19,962 for testing. All images were center- cropped and downsampled to 32 x 32 resolution with 3 RGB channels to reduce computational complexity while preserving key facial features.
As in our EMNIST experiments, we transformed each image into a pixel-level function over spatial coordinates. Specifically, pixel coordinates were scaled to the range [-0.5, 0.5) along both axes, and pixel intensity values in each RGB channel were normalized to the interval [-0.5, 0.5]. This setup enables the image completion task to be framed a conditional regression problem: the model is provided with subset of known pixel values (the context) and is tasked with predicting the remaining pixels (the targets).
Phtrps//www.nistt http:://wwwwnisstooooolplpooooee wwwff//pppaddttmmmeeeeeeeeeeeeea
hhtps://mmlab.i ie cuhk. edu. hh/prrojett//Celebb..ml
33

---

To simulate diverse conditioning scenarios, we sampled the number of context points lel ~ Unif(5, 45), and drew the number of target points It| Uni 50 50-|cl, ensuring that each training episode contains variable and realistic number of observed and queried pixels.
D Additional experiments
D.1 Ablation study
Ablation on training objective In this section, we conduct an ablation study on the training objectives used in our method. As described in $ 3.3, our full training loss consists of two components: (I) the KL di- vergence between two path measures, de- noted as LKI., and (2) the log-likelihood maximization loss, denoted as LLL- To understand the individual contribution each objective, we evaluate models trained with only one of the two loss terms.
Table 9: Log-likelihood results for the ablation study training objectives. In the second and third rows, the notation - A indicates that the objective was ex- cluded from the training loss during model training.
RBF
Model
context
target
TTSNP (ours)
0.893 #0.001
0.430 =0.001
-KL
0,880 #0.001
0,420
=0.001
CLL
0.872 #0.001
0,415 #0.001
As shown in Table 9, both LK and LL improve sample quality over the baseline variational posterior from the pre-trained model, highlighting their effectiveness enhancing the efficiency and diversity of particles in the learned SMCS procedure. However, the best performance is achieved when both objectives are combined, as TTSNP. This demonstrates that LKL and LL are complementary, working together to improve posterior approximation and overall inference quality.
SMCS
0.868 +0.001
0.405 #0.001
Number of samples In this subsection, following the setup in 3.1, we perform ablation study on the number of samples used during inference. As in other experiments described in $ 4, we include Pre-trained, Fine-tune, and SMCS as baselines for comparison. Both TTSNP and SMCS are evaluated with T' = 10 inference steps.
0.40
Log Liteiihood
0.35
0.30
Pru Lruised
SMGS
The experiment is conducted using a pre-trained NP model trained on 1D Gaussian Process data with an RBF kernel. For SMCS, apply the Unadjusted Langevin Algorithm as the transition kernel to sam- ple from the posterior during inference. The sults in Fig. 7 show that TTSNP achieves faster con- vergence compared other baselines as the num- ber of samples increases. This demonstrates that the learned intermediate distributions-constie using pseudo context representations-aale along with the learned forward and backward transitions, effec- tively guide latent samples toward the true posterior while capturing diverse modes, resulting in improved sample efficiency.
Fine Tuno
0.20
TISPP Goung)
25
100
200
1000
Number Samples
Figure 7: Comparison of target log- likelihood between TTSNP and the baselines 'Pre-trained', 'SMCS', and 'Fine-tune' us- ing NP model trained with GP data generated by RBF kernel. Here, we compare the log- likelihood under the same varying amounts of latent samples.
Ablation on the Model structure In this section, we conduct an ablation study the architectural design of the drift functions u and u. As detailed in Appendix B.4, our default implementation defines the drift functions follows:
u(ri, TT) = g²
(1-))) V log (r())) + (1 -1- 8.) NN(r., Doy Dp)),
(39)
2
o(r,, TT) =
log (Te) + (1 3r) NN(T,, t, D.. Dp)),
(40)
34

---

where NN(-) is implemented as a 3-layer MLP and is designed to approximate the score function V Vlog q(r D. U D,) of the intermediate distribution constructed from the context and pseudo context representations.
While this approach relies neural network approximate the full score function, an altemn- tive formulation can be derived by analytically decomposing the intermediate distribution using the Gaussian score form. In particular, we consider the following decomposition:
V log (r) 3 V p(r De) + (1 - V:) log 7(x (D))
(41)
+ (1 3,) log ( h(De, t, Ex) U UPp),
(42)
=V log म. (r) + (1 3.) r JUN (D.: Pp.t)
(43)
(D Dy.
where (x(r) denotes the part of the intermediate distribution excluding contributions from the pseudo context, and the neural network models the additional gradient signal induced by the pscudo context generator by modeling the mean JuNN and the covariance oNN of the Gaussian score. Also, more simply, we can design follows:
V log ग. (r) V log 7. (r) + (1 Be) PNN (De, Dp.
(44)
G²
where o² is the fixed variance, which we used for calibrating variational posteriors.
In this section, we compare the log-likelihood performance of TTSNP against the alternative drift modeling strategies described above, using ID) GP data generated with RBF kernel. All methods share the same pre-trained NP base model, allowing for a fair evaluation of how different formula- tions of the drift function impact inference quality. We ensure fair comparison by using the same training objectives, number of training samples, and number of inference steps across TTSNP and all other model variants-making the drift model structure the only varying factor. As shown in Table 10, while alternative neural network designs also improve sample quality and log-likelihood performance compared SMCS, TTSNP achieves the most significant improvement, highlighting the effectiveness of its architecture in modeling the score function and guiding posterior refinement.
D.2 Additional GP regression experiments
In Table 11, we extend the GP regres- sion experiments beyond those presented in 4.1 and $ 4.2 by exploring additional settings. Specifically, we first replicate the experiments originally conducted on ID GP data in & 4.1, but now using 2D GP data. This allows us empirically demon- strate that our method remains effective and directly applicable even when the in- put dimensionality increases. We first ex- tend the Sample Quality experiment de- scribed in $ 4.1 to more challenging set- ting using 2D GP data generated with an RBF kernel. The results are reported in the left table of Table 11. While the overall log-likelihood values are lower compared to the ID due to the increased com-
Table 10: Log-likelihood results for the ablation study on drift model structures. In the second and third rows, the notations 'Mean and Variance' and 'Mean' indicate that the neural network is used to model both the mean and variance of the Gaussian score function, and only the mean, respectively.
RBF
Model
context
target
TTSNP (ours)
0.893 #0.001
0.430 #0.001
Mean and Variance Mean
0.892 #0.002
0.420
#0.002
0.886 +.0000
0.421
+00002
SMCS
0.868 +-0.011 0.405 +0.001
plexity of the 2D input space- the performance trends remain consistent. Specifically, TTSNP continues to improve sample quality and enables efficient inference, confirming its effectiveness beyond simple ID tasks and into higher-dimensional scenarios.
We also replicate the Matern RQ experiment from $ 4.1 in the 2D setting. The corresponding results are shown in the right table of Table 11. Again, we observe similar performance patterns to those seen in the ID experiments, demonstrating that TTSNP enhances particle quality even in few-shot adaptation settings and generalizes effectively to unseen tasks. These results highlight the robustness of the learned intermediate distributions and transition kernels, even under high- dimensional covariate shift conditions.
35

---

Table 11: Loo-liikeihher results for Sample Quality (left) and Matern to RQ (right) scenarios.
Moddl
RBF
Model
Matern
RQ
contictt
contictt
contect
=0.110
1:01
Fine-tune
-0.156
1.0.005
0.384
20.007
Fine-tune
-0.202
200000
-0.440
20.005
0.310
10.004
0.092
1:0 005
SMCS
-0.154 10.000
00.667
20.0011
SMCS
-0.190 230.002
-0.432
20.002
0,318 20.003
0,104
200.00
TISNP -0.156 #0,004
40331
=0.005
TISNP
40,190
-0,420
0.316
0,110 #0.003
Table 12: Log-likelihood results on the CelebA task. The model was pre-trained on CelebA and evaluated after cross-domain adaptation.
Model
Context
Target
Pre-train
4.148110.00
2.028=000055
SMCS
4.15010.002
2.035=10.04
Fine-tune
4.138110.033
2.019-10.004
TTSNP
4.150=0.002
2.038=00.004
D.3 Additional Image Completion experiment results
Cross-Domain Generalization and CelebA Evaluation. Our image completion experiment was designed to evaluate the model's ability to generalize across domains. Specifically, we first pre- trained the DANP model on CelebA and subsequently trained only the transition kernel and pseudo representation generator on EMNIST, while keeping the parameters of the pre- pre-trained base model frozen. This setting isolates whether the model can adapt to EMNIST and corrupted EMNIST tasks without retraining the base network. Because the base model remains unchanged, its performance on CelebA under standard variational posterior sampling is expected to stay constant; hence, CelebA was not included in the main evaluation.
contribution of the SMC refinement procedure, allowing us to assess
Nevertheless, as shown in Figure 5, we include qualitative samples from both CelebA and EMNIST demonstrate that the model retains its generative capability on the source domain even after adap- tation. To further support this observation, we report the quantitative results on the CelebA dataset in Table 12. The results confirm that TTSNP maintains or slightly improves log- log-likelihood compared to the pre-trained and fine-tuned baselines, indicating strong robustness and generalization.
E Broader Impact
This paper identifies key limitation in latent Neural Process models: after pre-training, the global latent variables tend to be miscalibrated, resulting in low-quality samples from the variational pos- terior. To address this, we propose a method based on Sequential Monte Carlo sampler that refines these latent samples to better match the true posterior distribution. While this approach introduces additional computational cost at test time compared to standard inference with pre-trained models, the cost is adjustable based probabilistic inference quality. Aside from this computational trade-off, the proposed method poses no known negative societal impacts.
preferences and is justified by the significant improvement in
36


Reviews:
[{"review_id": "DUxKPAtVDs", "reviewer": "Reviewer_1HSn", "summary": "The authors identify an issue with NP: during inference, the learned posterior often deviates from the true posterior, especially at test time. This stems from the quality of the latent samples. Their proposal to address this issue is based on an SMC sampler operating in the latent space. This requires constructing a sequence of distributions, from initial to target. As a result, their proposal, termed TTSNP, \"transforms samples drawn from the variational posterior to be better aligned with the true posterior.\" The proposal is evaluated empirically.", "questions": "Please see my questions above", "limitations": "Partially covered.\n\nTo improve the paper, the authors could also comment on the practical application of the method, as the current experiments do not seem to reveal the complete potential of the developed methodology.", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "The paper is relatively clear, the idea is interesting, and the authors provide experimental evidence of the method's features. However, I fail to recognise the impact of the method, or what \"practical\" challenge (beyond the posterior mismatch problem mentioned) it addresses that no other method addresses in a better way. There are also some formatting/presentation issues in the manuscript.\n\nI give a more detailed list of points of improvement/questions: \n\n- Fig 3: What's the difference among all plots? This only seems to be explained in lines 270-276 and not in the fig's caption\n- It is unclear what the purpose of the experiments is. For instance, the proposed method is tested on image completion (the only experiment with real-world data). However, there are many other task-specific methods for image completion that -- according the results in fig 5 right --- could provide better results\n- The method is only tested empirically, mainly with synthetic data. Is it possible to provide any theoretical guarantee that the introduced SMC sampler in the latent space, improved the approximate posterior? \n- Captions are not properly produced: fig3 does not state what each plot is, fig2 is not referenced in the text, fig5 lacks \"left\" and \"right\" captions. \n- Why is the method referred to as \"scaling\" - what is \"scaled\" here? The method is said to \"sample more efficiently and effectively toward the true posterior\", but why is that to _scale_?", "was_revised": true, "final_justification": "Thanks for the thorough response. I appreciate the clarification on the points raised in my review, in particular the explanation about the appropriateness of the MCMC method considered and the experiments. After reading the rebuttal to my review (and that of other reviewers), I am increasing my score.", "rebuttal": "> **W1) Fail to recognise the impact of the method, or what “practical” challenge (beyond the posterior mismatch problem mentioned) it addresses that no other method addresses in a better way.**\n\nThank you for your insightful question. As you pointed out, the main contribution of TTSNP lies in addressing the discrepancy between the variational posterior and the true posterior of the latent distribution by employing Sequential Monte Carlo sampler at test time. This approach of refining variational posterior samples—originally underfitting the true posterior—through an exact sampling method to ensure better distribution is the first of its kind in the field of Neural Processes. As also acknowledged by other reviewers (qq3i, cBJS, afZz), this is an important and valuable direction for improving the quality and reliability of predictions when using NPs for statistical inference. Therefore, enhancing latent posterior sampling in neural processes models via TTSNP constitutes a significant methodological advance, enabling better inference in various practical regression and Bayesian optimization tasks.\n\n\n> **Q1) The purpose of the experiments is unclear. However, there are many other task-specific methods for image completion that could provide better results.**\n\nWe appreciate the reviewer’s attention to the rationale behind our experiments. The image completion experiment is designed to highlight the capability of NPs as uncertainty-aware meta-learning regression models to predict unseen parts of test images even with only a few given context points—sometimes fewer than 50 pixels in a 32×32×3 image. This task is a standard benchmark for evaluating NPs.\n\nIn our experiment, we demonstrate that, unlike other baselines, our model can successfully perform image completion on a target domain that was not meta-learned during the pre-training phase, by training only the transition kernel and the pseudo representation generator. Furthermore, we show that our method remains robust even when the context data is corrupted.\n\nWhile there certainly exist state-of-the-art methods and model architectures specialized for image completion tasks, our goal is not to outperform such methods. Rather, we aim to show that, similar to the GP regression case, refining latent samples using our method leads to clear performance improvement—underscoring the broader utility of our approach across different tasks that can be addressed by Neural Processes.\n\n> **Q2) Is it possible to provide any theoretical guarantee that the introduced SMC sampler in the latent space improves the approximate posterior?**\n\nThank you for your insightful and theoretically grounded question regarding the potential guarantees provided by the SMC sampler for improving the approximate posterior in the latent space. While our method is primarily motivated by empirical observations, there is indeed a theoretical foundation supporting the idea that **Sequential Monte Carlo (SMC) methods can improve posterior approximation**, especially when starting from an initial approximate distribution such as a variational posterior.\n\nAccording to theoretical results in the literature—particularly those discussed in [1,2]—SMC methods construct a sequence of intermediate distributions that gradually bridge the gap between a tractable initial distribution (e.g., the variational posterior) and the true posterior. By resampling and applying appropriate transition kernels, SMC ensures that samples are progressively guided toward the target distribution. Under certain regularity conditions, it has been shown that the **weighted particle approximation produced by SMC converges to the true posterior** as the number of particles.\n\nIn our case, the use of the variational posterior as the proposal for the initial distribution means that we start from a reasonable, albeit underdispersed, approximation. The SMC framework then allows us to refine this approximation by gradually correcting the mismatch through resampling and transition dynamics, effectively mitigating the underfitting issue. Therefore, while we do not derive a new theoretical result specific to our method, our approach is directly supported by established results in the SMC literature, and the **guarantee of asymptotic convergence** justifies its use for improving posterior quality in practice.\n\n**References** \\\n[1] Moral. et al. Sequential Monte Carlo samplers.\\\n[2] Chopin. Central limit theorem for sequential Monte Carlo methods and its application to Bayesian inference.\n\n> **Q3) Why is the method referred to as “scaling” - what is scaled here?**\n\nThank you for raising such a thoughtful question regarding the core of our work. *Test-time scaling* [1] refers to a class of methods in Large Language Models that enhance inference quality by increasing the *computational scale* at inference time. This includes **parallel scaling** techniques such as *self-consistency* [2] and *multi-agent decoding* [3], where multiple response paths are generated and aggregated, as well as **internal scaling** strategies like *Chain-of-Thought* prompting [4], where the depth or structure of a single response is expanded to improve reasoning.\n\nThese methods go beyond the standard greedy decoding of a single response path. Instead, they leverage multiple reasoning traces or longer generation paths to either combine multiple outputs or select the most plausible one—ultimately leading to improved inference quality.\n\nOur approach, TTSNP, shares this underlying philosophy: it increases computational effort at test time—either by using a larger number of SMC steps (T) or by sampling more latent variables—to refine inference. This refinement improves the quality of statistical inference similarly to how test-time scaling improves LLM output. This conceptual parallel is precisely why we adopt the term “*scaling*” in *Test Time Scaling for Neural Processes*.\n\n**References**\\\n[1] Zhang. et al. A Survey on Test-Time Scaling in Large Language Models: What, How, Where, and How Well?.\\\n[2] Wang. et al. Self-consistency improves chain of thought reasoning in language models. \\\n[3] Guo. et al. Large language model based multi-agents: A survey of progress and challenges.\\\n[4] Wei. et al. Chain-of-thought prompting elicits reasoning in large language models.\n\n> **Q4) Formatting issues regarding figure captions.**\n\nThank you for your valuable suggestion to improve the readability of the paper. First, we hereby provide additional information of each columns in Figure 3: (left) inference results using latent samples drawn from a variational posterior constructed solely from pseudo representations without access to the true context points; (middle) inference results using latent samples refined via  TTSNP’s SMCS procedure; and (right) inference results from the Fine-tune baseline. These figures demonstrate how well the generated pseudo representations cover diverse scenarios, and how the latent samples guided by such diversity can lead to better inference results compared to those drawn from the original variational posterior. Secondly, we will make sure to mention Figure 2, which is a schematic of our method, in Section 3. Lastly, we acknowledge that presenting Figure 5 in the same row as Table 3 may cause confusion; we will separate them in the final manuscript to improve clarity. In the final manuscript, we will follow your recommendation to enhance readability by improving the figure captions and reference quality accordingly."}, {"review_id": "MTjOHn8YnS", "reviewer": "Reviewer_qq3i", "summary": "The paper introduces a new method to enhance the performance of NPs by refining predictions at inference time using Sequential Monte Carlo and learned proposal transitions. The core idea is to improve posterior inference by resampling the global latent. The method uses generated pseudo-contexts to evaluate likelihoods and an SMC sampler learned at training time. \n\nThe authors conduct experiments across a variety of benchmark tasks and show that TTSNP enjoys superior performance and is broadly applicable.", "questions": "- I understand that the method is a general framework and can be incorporated with other methods. Have the authors conducted computational complexity comparisons with other methods? Is there a performance trade off for runtime? \n\n- How is resampling handled in the SMC algorithm? Is there adaptive resampling or effective sample size thresholds?", "limitations": "In addition to the weaknesses mentioned, have the authors tried their method on large underlying neural networks? What is the highest number of parameters that the experimented NNs have?", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\n\n- The paper is well-motivated; it addresses limitations like posterior underfitting and poor uncertainty calibration. In my understanding, the idea of refining NP posterior predictions at test-time via SMC sampling is novel within the NP literature.\n\nWeaknesses\n\n- It would be nice to have certain analysis on the approximated posterior. Discussions on how the SMC sampler should be used for different tasks (i.e. is convergence guaranteed?) would also supplement the material. \n\n- I am a little bit confused by the implementation of the learned transition kernels. It would benefit from more detail: are the transitions trained jointly with the base NP?\n\n- There is no ablation experiments on the number of pseudo contexts used. \n\n- I hope the authors could include a more comprehensive discussion on the “scalability” claimed in the paper.", "was_revised": true, "final_justification": "The authors have provided more explanation and justification that made me better understand the contribution of the paper. However, the authors could improve on their clarity by incorporating the explanation they've given in the rebuttal stage to the main paper.", "rebuttal": "> **W1) It would be nice to have certain analysis on the approximated posterior.**\n\nWhile our method is primarily motivated by empirical observations, there is indeed a theoretical foundation supporting the idea that Sequential Monte Carlo (SMC) methods can improve posterior approximation, especially when starting from an initial approximate distribution such as a variational posterior.\n\nAccording to theoretical results in the literature—particularly those discussed in [1,2]—SMC methods construct a sequence of intermediate distributions that gradually bridge the gap between a tractable initial distribution (e.g., the variational posterior) and the true posterior. By resampling and applying appropriate transition kernels, SMC ensures that samples are progressively guided toward the target distribution. Under certain regularity conditions, it has been shown that the weighted particle approximation produced by SMC converges to the true posterior as the number of particles and intermediate steps increases.\n\nIn our case, the use of the variational posterior as the proposal for the initial distribution means that we start from a reasonable, albeit underdispersed, approximation. The SMC framework then allows us to refine this approximation by gradually correcting the mismatch through resampling and transition dynamics, effectively mitigating the underfitting issue. Therefore, while we do not derive a new theoretical result specific to our method, our approach is directly supported by established results in the SMC literature, and the guarantee of asymptotic convergence justifies its use for improving posterior quality in practice.\n\n**References**\\\n[1] Moral. et al. Sequential Monte Carlo samplers.\\\n[2] Chopin. et al. Central limit theorem for sequential Monte Carlo methods and its application to Bayesian inference.\n\n> **W2) Implementation of the learned transition kernels. Are the transitions trained jointly with the base NP?**\n\nTo clarify, the transition kernels are trained separately from the base Neural Process model. Specifically, our experimental setting assumes that an NP model has already been pre-trained using an approximate ELBO objective. In practice, we empirically observed that the variational posterior of the pre-trained model underfits the true posterior distribution.\n\nTo address this, we introduce two additional components, namely the transition kernel and the pseudo representation generator, which are trained using only a small amount of additional data. Importantly, the parameters of the original NP model remain frozen throughout this training procedure. This ensures that performance improvements arise solely from refined posterior inference rather than additional training of the base model.\n\nFor a fair evaluation of our method, we also include a baseline called Fine-tune, where the same additional data is used to directly update the latent variational posterior through standard gradient-based fine-tuning. This comparison demonstrates the effectiveness of our approach in improving posterior quality without requiring full retraining of the underlying NP model.\n\n> **W3) Ablation on the number of pseudo contexts used.**\n\nThank you for your suggestion regarding an ablation study on the number of pseudo context points. In response to your comment, we conducted an ablation experiment on the GP regression with RBF kernel using the NP model, specifically examining how the number of pseudo representations affects performance. We use 12 pseudo representations by default and additionally evaluate configurations with 3 and 48. This comparison allowed us to assess how the model's performance varies with the quantity of pseudo context used for guiding latent inference. \n\n__Table R.1__ Ablation on the number of pseudo contexts.\n|Number of Representations | 3  |12(default)| 48    |\n|-----|----|----|----|\n|Context Log-likelihood |0.880±0.001|0.893±0.001|0.895±0.001|\n|Target Log-likelihood |0.423±0.002|0.430±0.001|0.438±0.001|\n\nThe results from this experiment will be included in the final manuscript to provide a more comprehensive understanding of the role of pseudo context size in our method’s effectiveness. Thank you again for raising this important point.\n\n\n> **W4) More comprehensive discussion on the scalability claimed in the paper.**\n\nThank you for raising such a thoughtful question regarding the core of our work. *Test-time scaling* [1] refers to a class of methods in Large Language Models that enhance inference quality by increasing the *computational scale* at inference time. This includes **parallel scaling** techniques such as *self-consistency* [2] and *multi-agent decoding* [3], where multiple response paths are generated and aggregated, as well as **internal scaling** strategies like *Chain-of-Thought* prompting [4], where the depth or structure of a single response is expanded to improve reasoning.\n\nThese methods go beyond the standard greedy decoding of a single response path. Instead, they leverage multiple reasoning traces or longer generation paths to either combine multiple outputs or select the most plausible one—ultimately leading to improved inference quality.\n\nOur approach, TTSNP, shares this underlying philosophy: it increases computational effort at test time—either by using a larger number of SMC steps (T) or by sampling more latent variables—to refine inference. This refinement improves the quality of statistical inference similarly to how test-time scaling improves LLM output. This conceptual parallel is precisely why we adopt the term “*scaling*” in *Test Time Scaling for Neural Processes*.\n\n**References**\\\n[1] Zhang. et al. A Survey on Test-Time Scaling in Large Language Models: What, How, Where, and How Well?.\\\n[2] Wang. et al. Self-consistency improves chain of thought reasoning in language models. \\\n[3] Guo. et al. Large language model based multi-agents: A survey of progress and challenges.\\\n[4] Wei. et al. Chain-of-thought prompting elicits reasoning in large language models.\n\n\n> **Q1) Computational complexity comparisons with other methods? Performance trade-off for runtime?**\n\nAs discussed in the Conclusion and in our response to W4, TTSNP is fundamentally a method that improves inference quality via computational scaling at test time. Consequently, it naturally incurs a higher computational cost compared to baseline methods that directly sample latent variables from the variational posterior and perform inference without refinement.\nTo quantify this increase in cost, we measured the actual computational overhead using the DANP model on the GP Regression task with an RTX 3090 GPU. The results, compared against the Fine-tune baseline under the same settings, are reported as follows:\n\n__Table R.2__ Time complexity comparison for 200 training steps.\n|   Method  | Fine-tune |  TTSNP|\n|------|-----|------|\n|Time (sec) |20 |80 |\n\nAdditionally, the relationship between test-time computational cost and inference quality is clearly visualized in Figure 4. As shown in the figure, increasing the number of SMC steps leads to a noticeable improvement in the quality of the latent variable particles. These particles become better aligned with the true posterior distribution, resulting in more accurate inference. Moreover, as the particle distribution converges more closely to the true posterior, we also observe a reduction in the variance of the log-likelihood, which indicates more stable and reliable predictions.\n\nThis empirical evidence supports the idea that TTSNP offers a favorable trade-off between computational cost and predictive performance, and that the degree of refinement can be flexibly controlled based on resource availability at test time.\n\n> **Q2) How is resampling handled in the SMC algorithm? Is there adaptive resampling or effective sample size thresholds?**\n\nThank you for your insightful question regarding the resampling step in our SMC sampler. As you rightly pointed out, this aspect was clearly described in Algorithm 1 in the appendix, but we acknowledge that it was not explicitly stated in the main text.\n\nTo clarify, we perform resampling when the effective sample size (ESS) falls below 30% of the total number of particles NN, i.e., when $ESS<0.3N$. This threshold is a conventional choice commonly used in SMC-based methods [1,2], and we adopted the same criterion in our implementation to ensure stability and sample diversity during the inference process.\n\nWe appreciate your suggestion and will make sure to incorporate this detail more clearly in the main body of the final manuscript, so that there is no ambiguity regarding the method design.\n\n**References**\\\n[1] Vargas et al. Transport meets variational inference: Controlled Monte Carlo diffusions.\\\n[2] Chen. Sequential Controlled Langevin Diffusions. \n\n> **L1) Highest number of parameters that the experimented NNs have?**\n\nThank you for your question regarding the scale of the underlying neural networks used in our experiments. We appreciate your interest in the model capacity and its impact on our method. In our experiments, the largest model we used contains approximately 1 million parameters, which is the largest parameter scale explored in the NP literature, consistent with the setting used in [1]. Also most of the experiments were conducted using the DANP model, which represents the most advanced architectural variant of Neural Processes to date.\n\nOur choice of this model and scale was intentional, as it allows us to evaluate the proposed method under a realistically challenging and expressive model, thereby demonstrating its applicability beyond small-scale or toy settings. While scaling to even larger architectures remains a promising direction, we believe that working with the largest and most expressive NP model to date provides strong empirical evidence of the effectiveness of our method within this research domain.\n\n**References**\\\n[1] Lee et al. Dimension agnostic neural processes."}, {"review_id": "mnnaNQJYsI", "reviewer": "Reviewer_cBJS", "summary": "The paper presents a novel approach for generating sample from meta-trained neural processes. The presented approach assume a sequential Monte Carlo approach for sampling the latent representation of the process conditioned on the context data. The variational posterior trained over the later representations is utilised as the proposal distribution $\\pi_0$ and the true posterior (which is known up to a constant) is used as the target distribution $\\pi_T$. The intermediate distribution $\\pi_t$ of the SMC is defined as annealed mixture of true posterior, variational posterior (but with variance fixed and set to $\\sigma^2=1$) and and variational posterior conditioned on the context data and a pseudo-context, which is essentially an artificial dataset generated by a neural network $h$. In order to train the forward and backward transition kernels needed for SMC, authors leverage a SDE where the forward and backward drift terms depend on an output of a neural network. This neural network is trained so that the marginal forward and backward measures at time $t$ coincide and are equal to $\\pi_t$. Authors evaluate the proposed approach on a number of synthetic and real-world problems.", "questions": "- For the experiments with a GP, can authors provide the likelihood of the analytical GP posterior?\n- In Section 4: Experiments, you write \"2) Fine-tune: the latent path is further adapted for inference by training the transition kernel using the available data at test time\", this sounds very vague. Can you provide more details on this procedure?\n- Can the approach of SMC with neural transition kernels be applied to an arbitrary Bayesian Inference problem?\n- Can you address the concerns I have regarding empirical evaluation (listed in the \"Weaknesses\" section)? For me this is the main issue with the current submission and I am willing to increase my score if these concerns are sufficiently addressed.", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths:\n+ The proposed method gives model the ability to correct itself during test-time, which is a conceptually novel approach that could potentially lead to very robust models\n+ The empirical results presented in the paper show the method delivers a consistent improvement over baseline, but I have some concerns about the evaluation, as listed in the \"Weaknesses\"\n\nWeaknesses:\n- Authors seem to be running the same experiments as in the DANP paper[31], but the results they obtain for the baseline model are vastly different, e.g. in Table 3 of DANP paper authors report the performance of 0.969 for EMNIST Target, but the submission in Table 3 reports 0.574 for EMNIST Target. There are no comments on where this significant discrepancy is coming from.\n- The method is only applicable to latent neural process. There is also a family of approaches that does not model the variational distribution over latent space, but instead models the distribution over the target variable directly that authors are aware of and cite [17,20,36] but do not compare against them in the evaluation. \n- Even though authors run the experiments on CelebA and show the completed images, they do not report numerical values. As such, I am not sure if there is much value in such an experiment, as we do not have any numerical measure of performance.\n- I believe a number of important ablations is omitted from the ablation study. For example, is the pseudo context set really necessary? Do we really need to replace the learnt variance of the variational distribution with $\\sigma=1$? These seem to be very critical questions that are left unanswered.", "was_revised": true, "final_justification": "In my review I raised a number of concerns regarding empirical evaluation and requested additional ablations. Authors sufficiently answered my questions and provided additional experiments, which resolved my most important concerns. I find the paper interesting and with issues with empirical evaluation resolved, I believe it should be accepted.", "rebuttal": "> **W1) Authors seem to running the same experiments as in the DANP paper, but the results they obtain for the baseline are vastly different.**\n\nThank you for your thoughtful comment and for your deep interest in our experiments. The difference between our results and those reported in [1] primarily arises from the differences in experimental settings.\n\nIn [1], the authors trained the DANP model jointly on both CelebA and EMNIST datasets, aiming to demonstrate the model’s capacity to handle data with varying dimensionality. This joint training exposes the model to both domains during meta-training, naturally resulting in strong performance across both.\n\nIn contrast, as detailed in our experimental setup, we pre-train the base model only on CelebA and then train the transition kernel and pseudo representation generator separately on the new dataset, EMNIST. The goal here is to assess robustness to domain shift by evaluating the model’s ability to adapt to unseen data without retraining the full model. Our method refines latent samples from the variational posterior in a post-hoc manner, using only limited additional data.\n\n**References**\\\n[1] Lee. et al. Dimension agnostic neural processes. \n\n> **W2) The method is only applicable to latent neural processes.**\n\nThank you for your question regarding the comparison with Conditional Neural Process (CNP) models. We focus primarily on Latent Neural Process (LNP) models for two main reasons.\n\nFirst, our experimental setup is designed specifically for LNPs. Our method targets the miscalibration of the variational posterior commonly observed in LNPs trained on large meta-training datasets. To address this, we introduce a test-time SMCS refinement procedure, training only the transition kernel and intermediate distribution networks while keeping the base model frozen. CNP models, lacking a latent path, cannot benefit from this refinement. Adapting a CNP model requires fine-tuning the encoder (and often the decoder), which risks overfitting and catastrophic forgetting, violating the frozen-base assumption that underlies our comparisons.\n\nSecond, prior work consistently shows that LNPs outperform CNPs under equivalent architectures. For example, Attentive Neural Processes outperform Conditional Attentive Neural Processes [1], largely due to the ability of latent variables to capture task-level uncertainty and global context, especially in low-data or extrapolation settings.\n\nThat said, in response to your request, we conducted additional experiments on the GP regression task including the CNP model, shown in Table R.3. For fair comparison, we fine-tuned both the encoder and decoder of CNP, unlike TTSNP and other baselines where only the latent path or auxiliary modules were trained.\n\n__Table R.3.__  Log likelihood comparison on GP regression task\n|Model | context| target|\n|---------|-----------|---------|\n|pre-train| 0.654±0.008 | 0.150±0.010 |\n|Fine-tune| 0.650±0.005 | 0.178±0.008 |\n|SMCS| **0.757**±0.004 | 0.186±0.006 |\n|TTSNP| 0.753±0.004 | **0.219**±0.006 |\n|CNP| 0.637±0.004 | 0.146±0.005 |\n|CNP_tune| 0.555±0.002 | 0.173±0.004 |\n\n**References**\\\n[1] Kim. et al. Attentive neural processes.\n\n> **W3) Numerical values for CelebA experiments.**\n\nTo clarify, our image completion experiment was designed to evaluate cross-domain generalization. Specifically, we first pre-trained the DANP model on CelebA and then trained only the transition kernel and pseudo representation generator on EMNIST, without modifying the parameters of the pre-trained base model. The goal of this setting was to assess whether the model could generalize to EMNIST and corrupted EMNIST tasks using only the SMC procedure. Since the base model remains frozen, its performance on CelebA using standard variational posterior samples remains unchanged. For this reason, CelebA performance was not included in the main evaluation.\n\nThat said, as shown in Figure 5, we presented samples from both CelebA and EMNIST to visually demonstrate that the model retains its generative ability on the original CelebA domain after adaptation. This was meant to illustrate that the learned components do not degrade performance on the source domain.\n\nWe agree that reporting CelebA performance can provide further evidence of robustness. Accordingly, we include the numerical results on CelebA in Table R.4 and will incorporate these quantitative results in the final version of the manuscript to support this point.\n\n__Table R.4__ Log likelihood results for CelebA task\n|Model | context| target|\n|---------|-----------|---------|\n|pre-train|4.148±0.001|2.028±0.005|\n|SMCS|4.150±0.002|2.035±0.004|\n|Fine-tune|4.138±0.003|2.019±0.004|\n|TTSNP|4.150±0.002|2.038±0.004|\n\n\n> **W4-1) Ablation study on pseudo context set.**\n\nThank you for your suggestion regarding an ablation study on the number of pseudo context points. In response to your comment, we conducted an ablation experiment on the GP regression with RBF kernel using the NP model, specifically examining how the number of pseudo representations affects performance.\n\nOur default setting generates 12 pseudo representations. To evaluate the impact of both increasing and decreasing this number, we additionally tested two alternative configurations in Table R.5: one with 3 pseudo representations and another with 48 pseudo representations. This comparison allowed us to assess how the model's performance varies with the quantity of pseudo context used for guiding latent inference. \n\n__Table R.5__ Ablation on the number of pseudo contexts.\n|Number of Representations | 3  |12(default)| 48    |\n|-----|----|----|----|\n|Context|0.880±0.001|0.893±0.001|0.895±0.001|\n|Target|0.423±0.002|0.430±0.001|0.438±0.001|\n\nThe results from this experiment will be included in the final manuscript to provide a more comprehensive understanding of the role of pseudo context size in our method’s effectiveness. Thank you again for raising this important point.\n\n> **W4-2) Ablation study on the replacement of the learnt variance of the variational distribution**\n\nThank you for your insightful comment regarding the ablation on recalibrated variance.\n\nThe recalibrated variance helps address the underestimation of uncertainty in the original variational posterior, allowing for better construction of intermediate distributions without additional learning. This leads to more accurate and stable inference during test-time refinement.\n\nTo validate its contribution, we conducted an ablation on the 1D GP regression task using DANP models, comparing performance with and without recalibrated variance in constructing intermediate distributions in Table R.6. This highlights the effectiveness of our variance adjustment strategy.\n\n__TableR.6.__ Ablation study on the variance\n|Type of Variance  |Learnt variance   |Fixed Variance|\n|----|-------|-------|\n|Context |1.272±0.001|1.273±0.001|\n|Target |0.463±0.003|0.661±0.004|\n\n> **Q1) Can authors provide the likelihood of the analytical GP posterior**\n\nIn our experiments with NP and DANP models, we used a GP regression setup with an RBF kernel, randomly selecting at least 5 context and 5 target points per batch, with a total cap of 50 points. The test set included 3000 batches of size 16.\n\nWe computed the true GP posterior analytically and evaluated log-likelihoods for both context and target points, obtaining averages of 1.3914 and 0.7265, respectively.\n\nTo evaluate TTSNP, we applied a DANP model pre-trained on 3D and 4D GP data to the 1D task without retraining. After test-time refinement, it achieved log-likelihoods of 1.2729 (context) and 0.6612 (target). These results show that TTSNP generalizes well to unseen domains and effectively refines the posterior, achieving near-optimal performance without additional training.\n\n> **Q2) Details of “Fine-tune”**\n\nAs outlined in our main experimental setup, we begin with an NP model pre-trained on a large dataset using an approximate ELBO objective. Following pre-training, we observe that the resulting variational posterior often deviates from the true posterior, particularly in few-shot or distribution-shifted settings.\n\nTo address this, we introduce a transition kernel and a pseudo representation generator trained using a small amount of additional data. These components refine the latent samples while keeping the pre-trained NP model frozen, ensuring that the core representation knowledge remains intact and enabling robust adaptation to new tasks.\n\nFor a fair comparison, the \"Fine-tune\" baseline uses the same additional data but updates only the latent path of the NP model, with all other components kept frozen. This setup ensures that both methods preserve the pre-trained representation while improving posterior quality through different mechanisms.\n\nWe agree that clarity on this configuration is important, and we will include a detailed explanation of both the \"Fine-tune\" baseline and our proposed training procedure in the appendix of the final manuscript.\n\n> **Q3) Can the approach of SMC with neural transition kernels be applied to an arbitrary Bayesian Inference problem?**\n\nOf course, it is possible. If we are performing Bayesian inference and already have access to a variational posterior over parameters or latent variables—whether it be learned through a training procedure or even based directly on the prior distribution—and assuming we can evaluate both the likelihood and the prior, then our method is applicable.\n\nIn fact, for an arbitrary Bayesian inference, one could even construct the framework using non-learned transition kernels such as those from the unadjusted Langevin algorithm. However, given sufficient data, one can still expect to construct a more effective and efficient transition kernel by leveraging a neural network. This learned kernel can better adapt to the structure of the posterior distribution and reduce sample inefficiency compared to fixed, hand-crafted transitions."}, {"review_id": "PJpKSVuAua", "reviewer": "Reviewer_afZz", "summary": "This paper proposes employing SMC samplers to improve the variational posterior in Neural Processes, with applications to meta learning. The authors claim that variational posteriors are typically miscalibrated, providing poor uncertainty quantification. \n\nThe authors propose a method that learns forward and backward kernels for the SMC sampler. The kernels are cleverly constructed to have desirable properties by construction, and are trained such that they are approximately time inversions of each other.\n\nThe kernels make use of additional synthetic data, which I understood to come from a generative model that is additionally trained. \n\nThe authors show experimentally that their method produces superior log-likelihood results.", "questions": "- How should we interpret the KL divergence  $KL(q( r | \\mathcal{D}_j) || q(r | \\mathcal{D}_c ))$? I read this as the divergence between the variational posterior conditioned only on the context in $j$ (the train data of task $\\tau_j$) and the variational posterior conditioned on the full data set for task $j$. I think I expected a divergence between some sort of prior and a conditional posterior\n- In Equation (2): does it ever make sense to use non diagonal covariance for $y$? \n- How is the prior $p(r)$ defined, and how does that relate to the claim of \"NPs can capture the data-generating mechanism underlying diverse tasks _without relying on explicit prior assumptions_\"? \n- Can the authors motivate using the metric of \"log-likelihood\" to measure the quality of uncertainty quantification?\n- What is the intuition for including log-likelihood in the loss for the transition kernels? Shouldn't only using the KL-loss lead to a more accurate posterior?\n- How is the model that produces pseudo contexts trained?", "limitations": "The authors have adequately discussed the limitations of their work.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "The proposed method is novel, and interesting. I also agree with the authors that uncertainty quantification and meta-learning are both very relevant fields of research. I like the motivation of these methods and while the experiments are not particularly extensive I still believe the results are significant enough to establish this method as an interesting area of future research.\n\nParts of the paper are very clearly written, with a strong story that is simple to follow because the authors often dedicate text to create intuition. I particularly like the background on SMC samplers and the second paragraph of section 3.2\n\nUnfortunately, other aspects of the paper are not so clear, even after consulting the appendix. For example, I do not fully understand how the model that generates pseudo contexts is learned. I also could not find how the prior $p(r)$ is defined.\n\nSome of the captions of Figures are also not very insightful. For example, Figure 3 is impossible to interpret without consulting the main text.\n\nFinally, the experiments do not fully support the claims made by the authors. For example, the paper mentions accurate uncertainty estimation, but most experiments test only for log-likelihood, which to the best of my knowledge tests predictive accuracy (even when tested on ood data.) For example, the authors make strong claims such as \"lack of diversity and poor calibration\" even though the experiment tests for predictive accuracy. \n\nI think the work is interesting but am leaning to reject this paper due to the clarity issues that I perceive. I would be happy to update my score appropriately if these issues are resolved.\n\nSome minor comments:\n- $\\tau_j$ seems to by synonymous with $j$, making the use of $\\tau$ for tasks unnecessary. Furthermore $\\tau$ is also used later for time in SMC, overloading the definition.\n- In Lines 73-74 and equations 2: While it is possible to infer the meaning from context, I think the functions $\\mu$, $\\sigma$ and $q$ could be introduced more properly.\n- Line 149: \"such as those arising from unseen kernels\" -> It is unclear what this means until the reader gets to the experiment section where GPs get mentioned for the first time.\n- The sentence in lines 21-24 is very long and difficult to parse\n- Line 25: have -> has\n- Line 48: I would write plural \"samplers\" instead of \"sampler\".\n- Line 51: \"with following\" -> \"with the following\"", "was_revised": true, "final_justification": "My original concerns were mostly in regards to presentation, and some confusion about the results and training procedure. I feel that the authors have adequately answered my questions and addressed the concerns of other reviewers. I am confident that they can incorporate their answers into the main text in time for the camera ready, as well as resolve the minor typos and errors.", "rebuttal": "> **W1, Q3) How is the prior p(r) defined? How does that relate to the claim of “NPs can capture the data-generating mechanism underlying diverse tasks without relying on explicit prior assumptions”?**\n\nThank you for your constructive comment regarding the prior distribution $p(r)$. To directly answer your question, when constructing intermediate distributions and the true posterior distribution within our SMC sampler, we adopted a standard Gaussian distribution as the prior. This choice facilitates efficient evaluation of the unnormalized densities required for importance weighting and transition updates during the SMC process.\n\nThe statement you referred to, \"NPs can capture the data-generating mechanism underlying diverse tasks without relying on explicit prior assumptions,\" relates closely to our discussion addressing Q1. As explained earlier, Neural Processes are trained to maximize the log-probability of target points rather than the full joint log-likelihood of context and target sets. Thus, the KL divergence in NP training is computed between variational distributions without explicitly modeling a prior distribution.\n\nHowever, in our framework, because our goal is to transform variational posterior samples toward the true posterior using SMC, evaluating the unnormalized posterior, which includes the prior, becomes necessary. To maintain tractability and ensure minimal informative bias, we utilize a standard Gaussian prior for latent variables. This allows efficient and effective sampling and weighting during the SMC procedure without imposing strong assumptions on the task structure.\n\n> **W1, Q6) How is the model that produces pseudo contexts trained?**\n\nThe pseudo context generator is trained jointly with the transition kernel using the same objective, composed of two main terms: the KL divergence between forward and backward path measures, and the marginal log-likelihood on target outputs.\n\nWhile the KL divergence ensures theoretically consistent latent dynamics, the log-likelihood term specifically guides the pseudo context generator toward producing diverse, task-relevant pseudo representations. By maximizing this log-likelihood, the generator encourages latent samples to capture multiple modes of the true posterior predictive distribution, thus increasing both the diversity and quality of generated samples.\n\nThis effect is clearly demonstrated in Figure 3 (left column), which shows that pseudo contexts generated by our approach lead to a broad range of posterior predictions. This diversity is crucial in enabling the transition kernel to refine latent samples toward accurate and uncertainty-aware predictions effectively.\n\n> **W2) Regarding the caption of Figure 3.**\n\nThank you for your valuable suggestion to improve the readability of the paper. Figure 3 illustrates the following across each column: (left) inference results using latent samples drawn from a variational posterior constructed solely from pseudo representations without access to true context points; (middle) inference results using latent samples refined through TTSNP’s SMCS procedure; and (right) inference results from the Fine-tune baseline. These figures demonstrate how well the generated pseudo representations cover diverse scenarios, and how the latent samples guided by such diversity can lead to better inference results compared to those drawn from the original variational posterior. In the final manuscript, we will follow your recommendation to enhance readability by improving the figure captions and reference quality accordingly.\n\n> **W3, Q4) Log-likelihood as uncertainty estimation?**\n\nThank you for your insightful question regarding log-likelihood as a metric for uncertainty quantification.\nLog-likelihood is widely recognized as a principled metric in Bayesian Neural Networks (BNNs) and Bayesian inference broadly, as it assesses both predictive accuracy and uncertainty calibration. Unlike deterministic metrics such as RMSE or MAE, log-likelihood evaluates the entire predictive distribution, making it especially suited to uncertainty-aware modeling [1, 2].\n\nHigher log-likelihood indicates that a model assigns greater probability density to observed data, reflecting predictive correctness and well-aligned uncertainty quantification. Theoretically, maximizing log-likelihood corresponds directly to minimizing the KL divergence between the model’s predictive distribution and the true data distribution, providing a statistically grounded evaluation of Bayesian inference quality [3, 4].\n\nIn our work, which leverages Sequential Monte Carlo methods to refine posterior inference in latent spaces, log-likelihood explicitly measures whether our refined latent samples produce predictive distributions that better approximate the true posterior predictive. Thus, log-likelihood serves as a rigorous metric to evaluate improvements in uncertainty-aware predictions, aligning directly with the goals of our method.\n\n**References**\\\n[1] Kendall. et al. What uncertainties do we need in bayesian deep learning for computer vision?. \\\n[2] Lakshminarayanan. et al. Simple and scalable predictive uncertainty estimation using deep ensembles.\\\n[3] Bishop. et al. Pattern recognition and machine learning.\\\n[4] MacKay. et al. Information theory, inference and learning algorithms.\n\n\n> **W4) Typos and minor comments.**\n\nThank you for your helpful comments regarding typos, equations, and phrasing. We appreciate your careful reading of the manuscript. We will incorporate all the suggested corrections into the final version to ensure the paper is clearer and more readable. Your feedback will certainly help improve the overall quality and presentation of our work.\n\n> **Q1) Form of the KL divergence?**\n\nThank you for your insightful question regarding the KL divergence term in the approximate ELBO for Latent Neural Process models.\nIndeed, the objective originally proposed in [1] typically involves a KL divergence between the variational posterior and the prior (as in Eq. 7). However, unlike the standard ELBO, which maximizes the joint likelihood of both context and target points, this modified formulation explicitly aims at improving predictive performance on target points at test time.\n\nThis adjustment leads to an intractable conditional prior distribution over latent variables given context points. To handle this intractability, the authors introduce a variational conditional prior as a practical approximation (Eq. 9 of [1]). Consequently, the resulting KL divergence in the training objective becomes a divergence between two variational distributions rather than between a variational posterior and the original prior.\n\nThus, the divergence between variational distributions naturally arises from approximating the intractable ELBO to specifically enhance predictive performance on target points, highlighting an important and subtle aspect of these training objectives.\n\n**References**\\\n[1] Garnelo. et al. Neural processes.\n\n> **Q2) Does it make sense to use non diagonal covariance of y in eq (2)?**\n\nIt is certainly possible in principle. However, models like DANP [1] are specifically designed to handle scenarios where both the input $x$ and output $y$ dimensions vary across tasks, enabling flexible inference across a wide range of tasks without requiring task-specific architectures. This flexibility is one of the key advantages of the model.\n\nIf one were to model the output variance using a non-diagonal covariance structure, additional design considerations would be required. Specifically, as the dimensionality of $y$ changes across tasks, the number of covariance parameters increases quadratically,  i.e., $O(n^2)$, which introduces computational and regularization challenges. Ensuring that the covariance matrix remains nonsingular and well-conditioned across tasks would also require careful design choices, such as appropriate parameterizations (e.g., low-rank plus diagonal structure) or regularization techniques.\n\nThat said, this direction offers a promising path toward more expressive and accurate uncertainty modeling, and we view it as a valuable avenue for future work. Thank you for your insightful comment regarding this potential extension.\n\n**References**\\\n[1] Lee. et al. Dimension agnostic neural processes.\n\n> **Q5) Intuition for including log-liklihood in the loss for the transition kernels? Shouldn’t only using the KL-loss lead to a more accurate posterior?**\n\nThank you for raising this insightful question. You're correct that our transition kernels are trained using both KL divergence and log-likelihood terms, each serving distinct roles.\n\nFirst, the KL divergence loss ($\\mathcal{L}_{\\text{KL}}$) between forward ($\\mathbb{P}^{u}$) and backward ($\\mathbb{P}^{v}$) path distributions ensures that our learned transition dynamics are theoretically consistent and time-reversible, aligned with recent score-based generative modeling approaches [1]. Minimizing this path-wise KL divergence guarantees optimal sampling properties but does not directly optimize for predictive accuracy or uncertainty quantification.\n\nSecond, the log-likelihood maximization term complements this by explicitly guiding the latent samples toward task-relevant predictions. Specifically, it encourages the generation of diverse latent samples that lead to accurate and uncertainty-aware posterior predictions.\n\nIn summary, the KL divergence loss ensures principled, theoretically sound latent dynamics, while the log-likelihood loss ensures practical effectiveness in predictive modeling. Together, they enable accurate inference with fewer particles.\n\n**References**\\\n[1] Chen. et al. Sequential Controlled Langevin Diffusions."}]

Author rebuttal:

Meta-review: This paper presents a neat take on the currently hot topic of "test time scaling" - i.e. spending more computation on model prediction to refine or improve samples from a model - applied to improve predictions from Neural Processes.  They develop a Sequential Monte Carlo sampler to refine initial samples from the variational posterior of Neural Processes towards samples from the true posterior.

The paper received 3 accepts and one borderline accept.  The reviewers initially seemed to find the presentation a bit confusing and in particular asked a variety of clarifying questions about the experiments.  However, the author's clarifications in the discussion seem to have been quite convincing as every reviewer ended up raising their score.  The reviewers found the paper novel, well written and the method compelling.  After discussion, they also found the experiments convincing.   Therefore, the recommendation is to accept the paper.
Decision: Accept (poster)