Adaptive teachers for amortized samplers
Authors: Minsu Kim
Sanghyeok Choi
Taeyoung Yun
Emmanuel Bengio
Leo Feng
Jarrid Rector-Brooks
Sungsoo Ahn
Jinkyoo Park
Nikolay Malkin
Yoshua Bengio
Venue: ICLR 2025 Poster (2025)
Keywords: amortized inference
generative models
reinforcement learning
GFlowNets
Primary area: probabilistic methods (Bayesian methods, variational inference, sampling, UQ, etc.)
TL;DR: We guide the training or amortized sequential samplers with a adaptive teacher model, leading to better mode coverage on a wide range of problems.

Abstract: Amortized inference is the task of training a parametric model, such as a neural network, to approximate a distribution with a given unnormalized density where exact sampling is intractable. When sampling is modeled as a sequential decision-making process, reinforcement learning (RL) methods, such as generative flow networks, can be used to train the sampling policy. Off-policy RL training facilitates the discovery of diverse, high-reward candidates, but existing methods still face challenges in efficient exploration. We propose to use an adaptive training distribution (the Teacher) to guide the training of the primary amortized sampler (the Student). The  Teacher, an auxiliary behavior model, is trained to sample high-loss regions of the Student and can generalize across unexplored modes, thereby enhancing mode coverage by providing an efficient training curriculum. We validate the effectiveness of this approach in a synthetic environment designed to present an exploration challenge, two diffusion-based sampling tasks, and four biochemical discovery tasks demonstrating its ability to improve sample efficiency and mode coverage. Source code is available at https://github.com/alstn12088/adaptive-teacher.

Published as a conference paper at ICLR 2025
ADAPTIVE TEACHERS FOR AMORTIZED SAMPLERS
Minsu Kim Mila, KAIST
Sanghyeok Choi KAIST
Taeyoung Yun KAIST
Emmanuel Bengio Recursion
Leo Feng Mila, Université de Montréal
Jarrid Rector-Brooks
Sungsoo Ahn KAIST
Jinkyoo Park KAIST
Mila, Université de Montréal
Nikolay Malkin University of Edinburgh
Yoshua Bengio
Mila, Université de Montréal
ABSTRACT
Amortized inference is the task of training a parametric model, such as a neu- ral network, to approximate a distribution with a given unnormalized density where exact sampling intractable. When sampling is implemented quential Decision-making process, reinforcement learning (RL) methods, such generative flow networks, can be used to train the sampling policy. Off-policy RL training facilitates the discovery of diverse, high-reward candidates, but ex- isting methods still face challenges in efficient exploration. We propose to use an adaptive training distribution (the Teacher) to guide the training of the pri- mary amortized sampler (the Student). The Teacher, auxiliary behavior model, is trained to sample high-loss regions of the Student and can generalize across unexplored modes, thereby enhancing mode coverage by providing efficient training curriculum. We validate the effectiveness of this approach in synthetic environment designed to present exploration challenge, two diffusion based sampling tasks, and four biochemical discovery tasks demonstrating its ability to improve sample efficiency and mode coverage. Source code is available at https:/ /etps://giit httgttttthuu.gg////////ggg.whh com/ als.n120stnn adaptive teacher.
a se-
INTRODUCTION
Sampling from complex distribution given its unnormalized density function is fundamental problem machine learning (Hinton, 2002; LeCun et al., 2006) and scientific discovery (Dellago et al., 1998; Noé et al., 2019). Amortized inference methods aim to fit a generative model that samples from a target distribution, possibly by sequence of stochastic generation steps which beneficial because it allows reusing a shared computational module for inference across multiple data points, as opposed to performing inference independently for each data point (Margossian & Blei, 2024). However, unlike for generative models trained from data, samples from the ground truth distribution may not be available. Multi-step sampling from an unnormalized density function with amortized inference can be achieved with reinforcement learning (RL) methods but raises the challenge of exploration - specifically, the ability discover new modes of the target distribution during training. This is due to the intractable size of the sample space and the fact that only sampling from the generator itself would be oblivious to modes that the generator misses.
Just as with generative models trained data, it is often more natural and beneficial to approx- imate the generation of objects as a sequence of decisions made by policy rather than using a single parametric family, due the multi-modal expressivity of hierarchical inference (e.g., diffu- sion probabilistic models (Ho et al., 2020)). Sequential decision algorithms for amortized inference unified by the theory of generative flow networks (GFlowNets; Bengio al., 2021), which a collection of off-policy RL methods (Tiapkin et al., 2024; Deleu et al., 2024). GFlowNets have been used for such amortized inference problems as natural-language and biological sequence de- sign by token-by-token sequence generation (Jain et al., 2022; Shen et al., 2023; Hu et al., 2024), Bayesian inference over data structures (Deleu et al., 2022), molecular design by incremental addi- tion atoms or fragments (Bengio et al., 2021; Jain et al., 2023), or image refinement by diffusion
"qual contribution, correspondence {mi n-sU, sanghyeok.. choi J@kai itraist.c. ac.ka kr

---

Published as a conference paper at ICLR 2025
Rehsvior palicy
Target distribution
Current distribution
Converged distribution
cooards
Student NN
Student
Student
Student
Tencher
loss valtes
tnisiing Unjocttiics
Teacher
Teacher NN
Roffer
espenences
Training an amortized sampler (Student) with an adaptive Teacher. Left: The hchavior policy mixes Student, Teacher, and replay buffer policies to generate trajectories that train Student and store experiences. Teacher is updated based on Student's 's loss. Right: Student and Teacher distributions co-evolve, with Teacher targeting uncovered modes until Student converges to the larget distribution.
Figure I:
process in continuous space (Venkatraman et al., 2024), inter alia. GFlowNets have shown success in the sequential sampling problems at scale due to their advantageous off off-policy training ability (Malkin et al., 2023).
The prudent selection of training data crucial to the success of such RL methods to model the full distribution faithfully, akin to the problem of active learning in supervised problems: to max- imize sample efficiency, the most informative samples should be selected for training. To explore the full distribution, some applications of GetswNets have either used exploration techniques bor- rowed from RL, such as noisy exploration (first used by Bengio et al., 2021), (prioritized) replay buffers (Deleu et al., 2022; Schaul, 2016; Vemgal et al., 2023), and delayed updates (Lau et al., 2023). Others have employed search techniques in the target space: MCMC and local search (Kim et al., 2024d:b; Sendera et al., 2024; Phillips & Cipcigan, 2024), genetic algorithms (Kim et al., 2024a), and exploiting samples from the target distribution (Zhang al., 2022; Hu al., 2023) when available.
However, challenges in mode coverage remain. All of the forementioned methods promote explo- ration through perturbation of the policy, e.g., replaying the samples (Vemgal et al., 2023), augment- ing the reward function (Pan et al., 2023b), and local search starting from generated samples (Kim et al., 2024d). These exploration methods focus capturing the modes that are already near those generated by the current policy and can hardly capture the ones sufficiently separated from the al- ready explored modes.
In this paper, we propose explicitly explore the regions of high loss by introducing Teacher model that guides the training of the primary, or Student sampler (Fig. 1). Here, we believe that trajectories with high loss are particularly informative for mode coverage, as they are likely to lead to regions of the target distribution that are either undersampled (dropped modes) or oversampled (collapsed modes). The Teacher is adaptive behavior policy that is trained to sample target space regions where the Student model receives a high loss. In turn, the Student model is trained on samples from the Teacher model.
Our approach can be seen as amortizing ideal prioritized experience replay (PER; Schaul, 2016) which samples high-loss objects from the entire sample space, instead of the finite-size replay buffer. Compared to (non-ideal) PER, the Teacher model has the potential to generalize across the high-loss regions of the Student, without regard for whether they have previously been sampled or discovered. In contrast, prioritized replay requires the poorly captured modes to have already been visited in order for the model learn from them.
We test our method on a diverse set of domains where GFlowNets have been used, including dis- crete tasks (biological sequence design and molecular design) and continuous tasks (diffusion-bar based sampling benchmarks). Comprehensive experiments demonstrate that our algorithm is effective in improving mode coverage and training efficiency across all tasks.
2 PRELIMINARIES
We give summary of GFlowNets as algorithms for amortized sampling by sequential decision making. For simplicity, this exposition about discrete space, where GFlowNets were originally defined (Bengio et al., 2021), but GFlowNets have been extended to the case of continuous variables (Lahlou al., 2023), which is conceptually similar. The reader is directed Bengio al. (2023) for an extended overview, Malkin et al. (2023) for an introduction focused on connections to hierarchical variational inference, and to Deleu et al. (2024) for maximum-entroppy RI. point of view.
2

---

Published as a conference paper at ICLR 2025
GFlowNets are policies in deterministic Markov decision processes (MDPs), trained as sample from a distribution over terminal states whose mass function proportional to a given reward. The MDP is assumed to be represented set of statey and (s s') E A if there is possible action to be taken at state s leading to state s'. A policy is then the same as collection of distributions Pr(- s) over the children' of every state s that has at least one child. As in other deep RI. methods, the policy could be a neural network (s(s' | s: ()) taking representation of the state as input and outputring the logits of a distribution over children s'. We will sometimes leave out the A as implicit to lighten notation.
finite directed acyelic graph G (S, where S is the
We assume the existence of unique state so E S, called the initial state that not the child of any state. Conversely, states without children are called terminal, and the set of terminal states is denoted X. A policy Py induces distribution over complete trajectories sequences = 5o - S1 Sn with S E X which can be sampled by starting at so and iteratively transitioning to child states sampled according Py until terminal state reached. This in turn terminating distribution P over X, which is the marginal distribution over the final states of trajectories sampled in this way. To be precise,
n-1
(T()) PF(T). PF(T (so - 5n))
Pr P I st),
(1)
where X and T indicates that the sum is restricted to trajectories whose last state is x.
Let R- X - R=0 be function on the terminal states, called the reward function, and set Z = Evex R(x) as the normalization constant. We would like to train Py so as to make (x R(x) for all x X, ie., to make Py sample terminal states with probability proportional to the reward.
Because the sum in (1) may be intractably large (if many trajectories could lead the same x), achieving this requires introducing auxiliary objects into the optimization. One popular option is use the trajectory balance (TB) objective (Malkin et al., 2022). To train model with TB, one introduces an additional backward policy Pa 1), which is a collection of distributions over the parents of every noninitial state (i.e., policy on the reverse MDP, which can be either fixed in our case - or learned), as well as an estimate of the total reward Z (usually parametrized the log domain, making log Ze learnable parameter). We define the TB discrepancy for trajectory 7 with final state x by
5(T; 0) := [log R(x) + log PB( [x)] [log Zo log (r(T:0),
backward flow
forward flow
where PB(T] Ix; 0) is defined analogously to (1), by
Pn(T= = (sn - ..... sn) xx)
Pn(s; 88+1). M 1=
(3)
It can be shown that if ō(T; 0 for all trajectories т, then Za Z and (x(x) = R(x)/Z for all x, meaning that Pp solves the sampling problem. Intuitively, this is the case because the reward distribution and Pa would then determine the same distribution over trajectories as Pp, but factorized in reverse order. One thus attempts to enforce this by minimizing a loss, such as ō(r; 0) on trajectories r sampled from some behaviour policy π. (If x is the current policy Pr itself, the optimization is said to be on-policy, otherwise off off-policy.)
We note that there exist other training procedures that use learned estimators and loss functions as- sociated with individual states or transitions, rather than full trajectories, such as detailed balance (DB; Bengio al., 2023) and subtrajectory balance (SubTB; Madan et al., 2023), which all have advantages under certain conditions. This paper mostly focuses on TB due to its simplicity and pop- ularity as a default choice in the literature; we investigate DB in Appendix F to show method's flexibility over objective functions.
If (s s') A, then s' is child of the converse relation is called parent.
3

---

Published as a conference paper at ICLR 2025
3 THE TEACHER: AN ADAPTIVE TRAINING DISTRIBUTION
In this section, we introduce the Teacher, which is a secondary GFlowNet designed to enhance the efficiency of off- poflic training for the primary, Student, GFlowNet. The two GFlowNets share the same state and action space, but have different rewards.
The Teacher's 's role is to generate an adaptive training distribution for the Student, aiming to sample trajectories that yield high loss for the Student. Intuitively, samples with high loss tend to be less (or never) visited by the Student, implying high probability the samples being in the unexplored modes. To this end, we train the Teacher with GFlowNet objective for amortization of sampling trajectories with high loss. Note that the Student's s target distribution does not depend on the Teacher, but the Teacher's 's target distribution depends on the Student.
We henceforth denote the parameters of the Student GFlowNet by 0 and those of the Teacher by s.
3.1 REWARD DESIGN FOR TEACHER
We define the reward function for the Teacher using the TB loss of the Student, 6(T;0)² In basic form, we could define the Teacher's 's reward
log R Teacher (x: (Pa(r(xe)) log (T::))))
(4)
In Eq. (4), the Student's loss is marginalized over trajectories 7 in the log domain over the backward policy (B(T x) of the Student, given terminal state This is because we aim to train the Teacher as a sampler of terminal states, although what we obtain when training the Student is a trajectory-level error ö(r; 0), which we need convert into function of the terminal state only to form the Teacher's ' reward. Having the Teacher model terminal states x rather than full trajectories is motivated by the desire obtain full mode coverage in the space of terminal states, but not necessarily in the space of trajectories that lead to these terminal states. In practice, this expectation is estimated using Monte Carlo sampling with single sample PB(r x: 8), relying on the fact that stochastic gradient descent training of the Teacher will automatically average out the variability resulting from this sampling. In fact, this gives unbiased estimator of the gradient if training with the full expectation (see, e.g., Deleu et al., 2022; Bengio et al., 2023).
We propose two modifications to (4) to facilitate mode discovery.
Favoring undersampled regions. First, we hypothesize that because the Teacher should encour- age the Student to discover unvisited modes, it should favor regions of the state space where the target density exceeds the Student's s sampling probability. To this end, we increase the weight of the Teacher's reward for states where the backward flow exceeds the forward flow (cf. (2)), while adding a smoothing constant E:
log
weighttdd
Teacher (x; (x; = (pa(()x:)
[oog (c+ (1+ Clacr: 8 (r: (T:
(S)
The weighted term (1+ + Clo(r:0)=0) gives additional weight when the TB discrepancy positive, which indicates the Student is undersampling righ-rewarded terminal sample. Here, C > 0 repre- sents the weighting constant, which we set to C = 19 for every task; see Appendix E.3 for ablation study on our choice of C.
Reward mixing. To ensure the Teacher covers the missing modes (the high-reward regions that the Student missed), it is important to focus the Teacher's search space more on the high-reward regions than the low-reward ones. This approach helps target both high-loss and high-reward areas effectively. To achieve this, we propose mix the reward (5), which is based the Student's loss, with the Student's 's log reward:
log Reeachrr (x) == log R weighttd
(x) log R(x).
This approach encourages the Teacher to sample regions with both high loss and high reward. Here the mixing constant a a hyperparameter that trades off between high loss and high reward; see Appendix E.4 for analysis of the effect of the choice of a.
4

---

Published as a conference paper at ICLR 2025
Algorithm Teacher-SSudent Training of GetowNets
1: Chuffer for r I....,Tu do
Initialize replay buffer with queue structure Iteration of training rounds
PF(r:0).
select Student
3: Select behavior policy PT(T)
PF(T:
select Teacher
PB(T)))P(x((uull Qhutte)))
select Prioritized Buffer.
4: Sample trajectories Tl..... TB PB(r).
Exploration
(Optional) Refine trajectories using local search.
6:
Compute rewards: R(x))...., R(xB)-
7:
Compute TB discrepancy of Student: 6(1):))...... 5(TB; 0).
8:
Compute Reacher (x1)...... Reacher using (R(x) ) and (ō(T;)))
9:
Compute TB discrepancy of Teacher: 6(T1: Φ), ..... 0(TB: d).
10:
Update Student parameters: Optimizer
1 Σ 0(T;
Student training
11: Update Teacher parameters: @ + Optimizer
1/2 Σ O(T; 2))
Teacher training
12: Add experiences to buffer: Qhulfer - Qulffer U (xy, R(x)). R Teachee (x,))
13:
end for
3.2 JOINTLY TRAINING TEACHER AND STUDENT
Using the (x; 0), the training process is joint optimization of the Teacher parameters & and the Student parameters to jointly minimize the following loss functions:
ZoPF(T;0) log P(x)
LSudent (T; 0) 6(T;0))
Teaachrr (T; () = 8Teacher (T; rr:
[log Reeacher (x; PB
ZoPF(T:
(8)
Notice that (7) is the loss for regular TB training of the Student, while (8) relies on the Student's loss provide reward for the Teacher via (6).
To simplify the training process, we adopt in experiments fixed backward policy PB(T | without trainable parameters, which used by both the Teacher and the Student.
xx),
Behavior policy for joint optimization. Algorithm I describes the Teacher-Student training pro- cedure, which simultaneously minimizes the loss functions of both the Teacher and the Student. In line 3, we select the behavior policy by choosing either the Student, the Teacher, or a prioritized buffer, ensuring that all three sufficiently utilized (see Appendix C for details on how they chosen). Given a terminal state x sampled from P(x Qbufter), we then generate a trajectory 7 us- ing the backward policy (B(T x). This approach is similar previous works (Shen et al., 2023; Sendera et al., 2024), which store only the terminal states x in the buffer and sample trajectories r using PB.
The behavior policy can produce an adaptive distribution of trajectories with respect to the Stu- dent's s learning state because the Teacher iteratively focuses on high-loss trajectories of the Student during training. This adaptivity is hypothesised to result in highly effective training for the Student.
Existence of stationary point of training process. The joint optimization for the parameters and over the support of (T) has a stationary point where the Student GFlowNet becomes an exact sampler and the Teacher GFlowNet samples proportional to eR(x)"; see Prop. 1 in Appendix B.
3.3 MITIGATING NON-STATIONARITY WITH LOCAL SEARCH
Joint optimization of the parameters 6, 0 with non-stationary target Reacher 0) poses significant challenges the Teacher's reward is nonstationary, evolving the Student learns. To address this issue, we use a local search method (Line 5) that locally optimizes (xeacee(x;; 0) based on the Teacher's samples. We expect the dynamic nature of Rxeacher 9) be effectively managed by such search, as the Teacher's 's main role is to generalize to modes poorly modeled by the Student, while the search helps the Teacher track the local changes in the Student's loss landscape.
Local search using a kernel defined by the policies first used by Zhang al. (2022); Hu et al. (2023) and extensively studied by Kim et al. (2024d)- involves iteratively backtracking trajectories and reconstructing them produce samples. The method consists of the following steps:
5

---

Published as a conference paper at ICLR 2025
1. Backtracking: Starting from terminal state x, we backtrack to an intermediate state using the backward policy Pg. denoted as (x....... -->
s).
2. Reconstruction: From the intermediate state s, we reconstruct a new terminal state x' using the
Teacher's forward policy Pr, represented as (s ..... .....'')).
3. Accept or reject: We accept the new sample x' in place of with acceptance probability A.
The acceptance probability A can be determined using either stochastic Metropolis- Netropolis-Hastinnss (MH) approach or deterministic ascent criterion (see Kim et al. (2024d) for details). This process is re- peated iteratively to progressively improve the samples so that their reward better matches (with MH) or locally maximizes (with the deterministic ascent version) the target reward function. UI- timately, we use the enhanced sample x' to train both the Teacher and the Student by generating trajectories ~ PB(T x²) and taking gradient steps on the losses (7) and (8).
4 RELATED WORK
GFlowNets. GFlowNets were originally introduced by Bengio et al. (2021) and extensively ex- tended by Bengio et al. (2023). They aim to develop sequential decision-making policy with a form of deep reinforcement learning that aims at sampling from the unnormalized density associ- ated with a positive reward function. Aiming to improve credit assignment over long trajectories, Malkin et al. (2022) introduced the trajectory balance (TB) objective mainly used in this paper. Building on this, Madan et al. (2023) introduced a mixing scheme that combines losses associated with subtrajectories, trading off the lower variance of DB with the lower bias of TB, Pan et al. (2023a) studied inductive biases that use partial reward information, and Jang et al. (2024h) ex- tended this idea to learnable reward shaping schemes. Shen et al. (2023) and Jang et al. (2024a) studied auxiliary losses for better training of the backward policy.
Orthogonal to those studies, other works focus on improving off-policy training. Deleu et al. (2022); Shen et al. (2023); Vemgal al. (2023) studied the use of replay buffers in GFlowNets to enhance sample efficiency. Kim et al. (2024d;a;b); Sendera et al. (2024) investigated local search methods to guide GFlowNets toward high-reward regions. Kim et al. (2024c) proposed to adjust the exploration- exploitation trade-off via amortized conditioning on reward temperature. Similarly, Lau et al. (2024) introduced method that mixes Deep Q-Networks (Mnih, 2013) (exploitation) with GFlowNets (exploration) to balance the exploration-exploiiation trade-off. Our proposed method is also an off- policy training approach for GFlowNets. In contrast to the methods above, which use off- training to focus GFlowNets on high-reward regions, our method aims to address missing modes and underexplored regions. Note that the aforementioned reward reward-seeking off-police methods are complementary our approach; for example, local search with Teacher is studied in $5.1.
While the above algorithmic work mostly concerns discrete space, Lahlou et al. (2023) intro- duced the theory of GFlowNets in continuous space, leading to subsequent work on diffusion sam- plers (Zhang et al., 2024; Sendera et al., 2024), posterior sampling under diffusion priors (Venka- traman et al., 2024), and applications to molecular dynamics (Scong et al., 2024). Our proposed algorithms are effective in both discrete and continuous space ($5.2).
Adaptive training distributions. Adaptive training distributions are essential techniques in deep learning, ensuring that the training data evolves appropriately during model training. For instance, curriculum learning methods (Bengio al., 2009) schedule the difficulty of training tasks by gradu- ally increasing from easy to hard, thereby facilitating more efficient model training. These methods are also widely applied in reinforcement learning, e.g., Narvekar al. (2020).
Active learning (Gal et al., 2017), which usually built for supervised leaming, also falls into this category. where the training dataset actively changes to discover better strategies. Especially, uncertainty sampling-based methods (Sener & Savarese, 2018; Yoo & Kweon, 2019; Kirsch al., 2019; Ash et al., 2020) that prioritize sampling data points having high predictive uncertainty of the classifier are relevant to our idea. The key difference is that our method is built for reinforcement learning, where we do not rely on the predictive uncertainty of a classifier but on the loss value defined by the compositional policy.
Our work is also relevant to few-shot experimental design (Wang et al., 2024), our Teacher plays the similar role as the
adversary that generates tasks.

---

Published as a conference paper at ICLR 2025
Table I: Evaluation results on deceptive grid worlds with dimension d and grid length H. The number of modes discovered (# modes) and empirical L1 distance between target and sampled distributions are reported as =standard deviation over five runs. The L1 distances are scaled appropriately for readability.
Grid contig. ->
2, 11 - 128
d 2, 256
d- 4, HI
d-
Algoriihm Metric
modes (T)
L, 10 P()
#modes L, 10 (L)
modes (T)
L, x $ ())
modes
10
(on-policy
2
1.74
0002
16.6=
1,635
e-cxpl.
555.2 66.2
3.59 + 0.66
672.6 16.3
1,75 +
6.6+1.2
1,000 +
24.2+
16634 0.000
GAFNN
675.4 0.5
1.66 0.11
1044,8 276.
1,55 +.17
11.8
1,57 0.002
24.6=74
1 1,664 + 0.022
676.0
1.55
a
1.648
669.0 3.8
4.88 0:44
2055.2
1.71 0.08
16.0+ 2.8
+ 0.034
46.6 14.6
1,639 0.01
Teacher lowrs)
676.0 a: 0.0
2.13 0.18
2452.6 21.7
0.94 0.03
$1.4 4.0
1.019 =: 0.016
246.6 14.7
1.634 0.001
Target distr.
Teacher (ours)
PER
PRT
GAFN
E-expl.
On policy
Figure 2: Empirical distribution plots of 10 test samples from policies on the (d 2, 256) grid.
5 EXPERIMENTS
This section provides empirical validation of our method. Our primary goal is to demonstrate that approach is beneficial over other off-policy training methods for trajectory balance (TB), repre- sentative learning objective for amortized samplers. We also aim to show that our method can effec- tively be integrated with existing off- off-policy search methods, such local search and replay buffer techniques. We benchmark our approach on three major tasks: two discrete tasks ($5.1 and 5.3) and one continuous task ($5.2). We also include three ablation studies and pilot experiment on inte- grating local search in Appendix E. Additionally, we tested the versatility of our method by adopting another training objective, detailed balance (DB; Bengio et al., 2023) in Appendix F.
5.1 DECEPTIVE GRID WORLD
Setting. The deceptive grid world is a synthetic environment modified from the grid task intro- duced by Bengio et al. (2021). It consists of a d'dimensional hypercube of side length H. resulting in a search space of size O(H²). The agent starts from the origin (position 0) and can only move in directions that increase coordinate by 1 or terminate to receive the reward. The reward of each terminal state x = (X),..... (x).....xa) is given by
d
d
R(x) = Ro + Ri
1
x
0.5 < 0.1
+ R₂
N i=1 H-1 o.sle (0.3.0.4)
(9)
H-1
The modes with the highest rewards (R₂ + Ro) are surrounded by walls with low rewards (Ro). In between them are deceptive regions offering relatively high rewards (Ri + Ro), which can lure the agent into getting trapped. We Rg = 10-5, 5, R₁ = 0.1, and R₂ 2. Following previous works, we use the number of modes discovered and the empirical L1 distance from the target distribution as evaluation metries. See Appendix D. for more details about the settings.
Baselines. We compare our method with on-policy TB and TB with off-policy exploration meth- ods, such as e-cxploration (Bengio et al., 2021), and GAFN (Pan et al., 2023b), along with baseline using a replay buffer prioritizing rewards (PRT) or Teacher rewards inspired by Prioritized Experi- ence Replay (PER; Schaul, 2016). PER can be seen as a non-amortized version of our method.
Results. Table 1 summarizes the results. TB with Teacher consistently outperforms baselines. The significant margin in the number of modes discovered the larger-scale setting indicates that the Teacher effectively guides the Student to visit undiscovered modes. Please refer to Appendix E. for comprehensive results.
Effects of local search. We assess the local search (LS) effect ($3.3) on (d 4,H 32) grid. The Teacher's local search uses the Teacher's reward Reacher for acceptance, compared to two baselines on-policy TB and TB with PER which the Student's reward R. Furthermore, local search accelerates mode discovery, likely by reducing the sensitivity to nonstationarity in Teacher learning (See Appendix E.5).
7

---

Published as a conference paper at ICLR 2025
Table 2: Evaluation on multimodal continuous sampling tasks. Log-partition function estimation errors (ev- idence lower bound (ELBO), importance sampled ELBO (ELBO-IS), evidence upper bound (EUBO)) and 2- Wasserstein distances (W2) to target samples are reported as mean+sd over five runs. We compare MCMC methods (SMC, GGNS), differentiable simulation method (PIS), and GFlowNets trained using the TB objec- tive with various off policy strategies, such loss prioritized replay (PER) and reward prioritized replay (PRT).
Energy
25GMM (d 2, log = 0. W3 = 0.29)
Manywell (d 32, log Z 164.6956733, W3 = 5.36)
Algorithan Metric-
ELBO (T)
ELBO-1S (1)
EUBO (L)
W (1)
ELBO (1)
ELBO-IS (T)
EUBOX))
W3 (1)
SMC
0.56922006
0.86
149.70621.00
8. 8.288t.32
GGNS
0.016x2042
19
164.404z0 ase
6.51
PIS
-1.192
-1. 192
26.733-5.00
4.95
160,516a ocs
162.01 7
381.464-200.00 916
6. 150001
TB
(on-policy -]
-10008
18.321
4.64.m01
162.019musss 161.048
407.850.uoos:
6.15
e-cxpl.
-1,056un.117
-0.9 10
15.135-1.01
4.58-a06
161 161.06422066
162. 008
355.787=47i
6.15-470
4.7500mt
0..640au444
12.103.221
4.28ausv
161.07
161, 91
379.623.m.4s
6.14.007
PER
-0.282an.15 i58
-0.147anu18
18833=2166
188a112
161.537
162.582-n2s6
210.4-10aasss
sse
5.91
Teacher (owss)
-0.137=0000
-0.005
0.115=000e
0.86
163.484:004p
164.676=n.04s
165.800
5.462001
Target distr.
Teacher (ours)
PER
PRT
E-expl.
On policy
Figure 3: Samples from trained models the Manywell task (projected onto the first two dimensions).
Goal distr.
Student (1/5)
Teacher (1/5)
Goal distr.
Student (2/5)
Teacher (2/5)
Figure 4: KDE plots for 25GMM (left three) and Manywell (right three) at intermediate states of training. The Student (ratio) indicates the fraction of total training steps completed. The Teacher adaptively adjusts the training distribution in response the modes that the Student missing.
5.2 DIFFUSION SAMPLING
Setting. The task of learning diffisson sampler is to invert a diffusion process in order to sample the target density function. These experiments largely follow the setup of Sendera al. (2024). Here, we aim to model a distribution over trajectories (with As = 4: for us, T = 100), so that x1 is distributed according to an unnormalized density function p(x)) R(x,) = e -E(x)) The transitions are parametrized as to the Guler-Maruyama dis- cretization of a neural stochastic differential equation (Tzen & Raginsky, 2019): the sampler begins at the initial state (0.r r = 0). and the transition from to Xrtar is sampled from an appropriately scaled Gaussian with mean given by trained model taking X, and as input. The detailed setting and policy parametrization are described in Appendix D.2.
(0 = - XAr + X221 - - x1)
The main challenge in this task to capture the multimodality of R(x) without having access to sam- ples from target distribution during training, which difficult when there are many well-separated modes. It is not possible to apply the forward KL (i.e., log- likelihhhoodd variational bound) objectives typically used for diffusion models (Song et al., 2021), since target distribution datapoints are not available; thus exploration is crucial to mode discovery.
In this work, we benchmark diffusion samplers on two established tasks in the diffusion sampling literature: 2-dimensional Gaussian mixture with 25 modes (25GMM) and 32-dimensional Many- well distribution. When performing off-policy exploration, we assume a black-box property for the energy function 8, where VE is not accessible, similar settings in reinforcement learning. This as- sumption is meant to mimic a common setting in scientific discovery, where we may have black-box energies requiring expensive simulations compute; note that for these tasks, effective methods that use the energy gradient exist (see Sendera et al. (2024)).
Baselines. We consider two representative MCMC baselines: Sequential Monte Carlo (SMC) and the state-of-the-art GGNS (Lemos al., 2024). We also include the simulation- based continuous
8

---

Published as a conference paper at ICLR 2025
QMP (")BO
QMP (BUBO (13)
QM9 (Num Modes ("3)
25
2.3
15
-2.5
10
200
-30
0
200
603
1000
200
-00
600
500
1000
200
500
500
1000
Nam.. Training Rounds
Narn. Traliing Rournds
Nam TToiring Rourds
sEH (113
sEH (EUBO (U3)
SEH (Num Modes (?1)
100
30
-30
22
-10
1000
1901
330
2000
933
1880
1980
Sun. Truniin Rounds
Nam. Training Rounds
Nun. Traning Rounds
ThhindX (ELBO
TFFinds (EUBO
TThinci (Num Modes (13)
300
12
00
100
-25
10
-001
600
1000
600
1000
Nam. Trainig Rounds
Nari. Trainitg Rounds
Nast. Trainit Rounds
L.1-RNAI (T)
L11-RNAI (EUBO
1.14 RNAI (Num Modes (T))
1000
750
-10
2:0
1008
1900
2000
530
1000
Sun. Traiting Rounds
Nam. Traising Round
Nun. Traning Rousds
TB PKT
18 Tencher PKC
TB + E-capl.
TB PER
TB- Teacharr PER
Figure 5: Training graphs for molecule design (QM9, sEH) and biological sequence design (TFhind8, L14-RNAI) tasks. Mean and standard deviation over five runs are shown.
stochastic control method Path Integral Sampler (PIS; Zhang & Chen, 2022). The major base- lines are off-policy training methods based on trajectory balance (TB). We include on-policy TB in continuous space and the e-exploration method introduced by Malkin al. (2023); Lahlou et al. (2023), which adds additional Gaussian noise at least policy sampling step during training, as main baselines. Additionally, prioritizes reward (PRT) as studied by Sendera et al. (2024), and another that prioritizes loss (PER) (Schaul, 2016). The gradient-based local search introduced by Sendera et al. (2024), while effective, excluded as requires access to Vε(x) for the search; but we also study potential integration of the Teacher with these techniques in Appendix E.5.
compare with two replay buffer methods combined with TB: one that
Results. As shown in Table 2, on-policy TB and PIS yield similar performance, consistent with the fact that they have identical expected gradients (Malkin et al., 2023; Lahlou et al., 2023). This suggests that TB could benefit from additional off-policy methods for improvement. Indeed, the e-exploration techniques improve slightly over on-policy methods. PRT provides larger benefits on 25GMM but shows meaningful benefits on Manywell. PER offers significant improvements over on-policy methods compared to others. Our method, Teacher, achieves the highest results across all metrics, including the Evidence Lower Bound (ELBO), Importance-Sampe SSmpleed ELBO (ELBO-IS), Evidence Upper Bound (EUBO) (see Appendix for detailed definitions of those metrics). Espe- cially in EUBO, a metric suggested by Blessing et al. (2024) to measure mode coverage, baseline methods face significant challenges. This indicates that existing methods struggle to perform proper exploration across modes, as confirmed by the sample plots in Fig. 3. Our method offers clear advantages on EUBO metries.
In Fig. 4, we depict the training dynamics of the Teacher and Student by plotting kernel density estimates of their samples midway through training. This figure illustrates the mechanism by which the Teacher promotes mode discovery by the Student. As the Student struggles to find some modes, especially those with lower rewards than others, the Teacher puts high probability on them, encour- aging the Student to reduce its loss in those regions of the space.

---

Published as a conference paper at ICLR 2025
5.3 BIOLOGICAL AND CHEMICAL DISCOVERY
Setting. GetowNets have been used to generate biological and chemical structures by sequentially adding predefined substructures. In molecules, the actions add atoms or fragments; in biological sequences. nucleotides or amino acids. We aim to match a target probability distribution over struc- tures given by some proxy reward model. Following Shen et al. (2023), we benchmark the number of discovered modes, as well as probabilistic metries like ELBO and EUBO. We study four biological and chemical discovery problems, following Shen et al. (2023); Kim et al. (2024d):
QM9. The objects being sampled are small molecular graphs. Molecules are generated using 12 building blocks with stems, and each molecule contains 5 blocks. The reward function is HOMO-LUMO gap on the target transcription factor, which is obtained via a pre-trained MXMNet proxy from Zhang al. (2020). We use reward exponent of 5. We define modes as the top 0.5% quantile of R(x).
sEH. The generated objects are molecular graphs. Molecules are built using 18 blocks with stems and 6 blocks per molecule. The reward function is a binding affinity to soluble epoxide hydrolase (sEH), which is provided by the pre-trate pre-trained proxy model from Bengio et al. (2021). We use a reward exponent of 6. We define modes as the top 0.01% 0 quantile of R(x), with filtering to exclude candidates too similar based on Tanimoto similarity, following Kim et al. (2024c).
TFbind8. The generated objects are DNA sequences with nucleotides. The reward function is a binding affinity to a human transcription factor (Barrera et al., 2016), which is obtained via pre-trained proxy model provided by Trabucco et al. (2022). We use a reward exponent of 3. We use a pre-defined set of modes provided by Shen et al. (2023).
L14-RNAI. The generated objects are RNA sequences of length 14. The reward function is a binding affinity to human transcription factor, which is obtained via pre-trained proxy model from Sinai al. (2020). We use a reward exponent of 8. We define modes as the top 0.01% quantile of R(x), with diversity filtering whose threshold is 1 unit of Levenstein distance, also following Kim et al. (2024c).
Detailed training and hyperparameter settings for each task can be found in Appendix D.3.
Baselines. We compare our method with on-policy TB, e-exploration, and the PRT replay buffer method designed for biochemical tasks (Shen et al., 2023). Additionally, we evaluate it against a loss-prioritized replay buffer (PER) (Schaul, 2016). Since the local search method (Kim et al., 2024d) targets reward exploitation with different purpose to ours, a separate analysis of its com- plementarity to our method presented in Appendix E.5.
Results. As shown in Fig. 5, the Teacher method improves mode discovery when combined with both PER and PRT buffers, outperforming on-policy methods every task. Similar trends are observed across other metrics, with faster convergence under the Teacher method. For TFbind8, our method's s dominance is particularly evident for mode coverage metrics like EUBO and the number of modes, although ELBO remains comparable using PER or PRT. In the largest task, L14-RNAI, the Teacher method surpasses PER and PRT in ELBO, EUBO, and the number of modes.
We draw special attention to the comparison between PER and PRT. While the differences are min- imal in the first three tasks (QM9, TFbind8, sEH), in the larger-scale tasks, pessorrroriiizing with PER shows a clear advantage. This suggests that loss information is more important for exploration in large-scale tasks, where the Teacher method with PER achieves the best overall performance.
DISCUSSION
We have introduced Teacher that adaptively generates states for training Student amortized sampler. The Teacher favors states for which the Student has high loss and thus promotes the discovery of new modes.
This approach paves the way for numerous future research directions. For example, an adaptive Teacher could be applied to amortize intractable inference in large language models (LLMs) (Hu al., 2024) and diffusion models (Venkatraman et al., 2024), tomatic red-teaming for LLMs (Lee et al., 2024). Applying our method to probabilistic models such amortized Bayesian causal discovery (Deleu et al., 2022; 2023; Kasikawa-Toooee 2022) and amortized inference in graphical models (Falet et al., 2024) are also promising directions. Methodologically, the concept of using amortized prioritized experience replay (PER) as Teacher can be extended to train general agent-based systems, not only amortized samplers.
enhance exploration in au-
al.,
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGEMENT
We thank to Anirudh Goyal. David Dobre, Gauthier Gidel for helpful discussion for this project. This research is based funding from Samsung, Intel, CIFAR and the CIFAR AI Chair pro- gram. The research was enabled in part by computational resources provided by the Digital Re- search Alliance of Canada (https: ////a/alll allianneccan ca), Mila (https: / //mi la quebec), and NVIDIA. This work was also partially supported by the National Research Foundation of Korea (NRF) grant funded by the Korea government (MSIT) (No. RS-2024-00040008))
REFERENCES
Jordan T Ash, Chicheng Zhang, Akshay Krishnamurthy, John Langford, and Alekh Agarwal. Deep batch active learning by diverse, uncertain gradient lower bounds. International Conference on Learning Representations (ICLR), 2020.
Luis A Barrera, Anastasia Vedenko, Jesse V Kurland, Julia M Rogers, Stephen S Gisselbrecht, Elizabeth J Rossin, Jaie Woodard, Luca Mariani, Kian Hong Kock, Sachi Inukai, al. Survey of variation in human transcription factors reveals prevalent dna binding changes. Science, 351 (6280): 4504-454, 2016.
Emmanuel Bengio, Moksh Jain, Maksym Korablyov, Doina Precup, and Yoshua Bengio. Flow net- work based generative models for non-iteratin diverse candidate generation. Neural Information Processing Systems (NeurIPS), 2021.
Yoshua Bengio, Jérôme Louradour, Ronan Collobert, and Jason Weston. Curriculum learning. In- ternational Conference on Machine Learning (ICML), 2009.
Yoshua Bengio, Salem Lahlou, Tristan Deleu, Edward J. Hu, Mo Tiwari, and Emmanuel Bengio. GFlowNet foundations. Journal of Machine Learning Research, 24(210):1-55, 2023.
Denis Blessing, Xiaogang Jia, Johannes Esslinger, Francisco Vargas, and Gerhard Neumann. Be- yond ELBOs: large-scale evaluation of variational methods for sampling. arXiv preprint arXiv:2406.07423, 2024.
Xinyue Chen, Che Wang, Zijian Zhou, and Keith W. Ross. Randomized ensembled double Q- learning: Learning fast without a model. In International Conference on Learming Representa- fions (ICLR), 2021.
Tristan Deleu, António Góis, Chris Emezue, Mansi Rankawat, Simon Lacoste-Julien, Stefan Bauer, and Yoshua Bengio. Bayesian structure learning with generative flow networks. Uncertainty in Artificial Intelligence (UAI), 2022.
Tristan Deleu, Mizu Nishikawa- Toomey, Jithendaraa Subramanian, Nikolay Malkin, Laurent Char- lin, and Yoshua Bengio. Joint Bayesian inference of graphical structure and parameters with single generative flow network. Neural Information Processing Systems (NeurIPS), 2023.
Tristan Deleu, Padideh Nouri, Nikolay Malkin, Doina Precup, and Yoshua Bengio. Discrete prob- abilistic inference as control in multi-path environments. Uncertainty Artificial Intelligence (UAI), 2024.
Christoph Dellago, Peter G Bolhuis, and David Chandler. Efficient transition path sampling: Ap- plication to lennard-jones cluster rearrangements. The Journal of chemical physics, 108(22): 9236-9245, 1998.
Pierluca D'Oro, Max Schwarzer, Evgenii Nikishin, Pierre-Luc Bacon, Mare G Bellemare, and Aaron Courville. Sample-efficient reinforcement learning by breaking the replay ratio barrier. International Conference on Learning Representations (ICLR), 2023.
Lean-Pierre René Falet, Hae Beom Lee, Nikolay Malkin, Chen Sun, Dragos Secrieru, Dinghuai Zhang, Guillaume Lajoie, and Yoshua Bengio. Delta-Al: Local objectives for amortized inference in sparse graphical models. International Conference on Learning Representations (ICLR), 2024.
Yarin Gal, Riashat Islam, and Zoubin Ghahramani. Deep Bayesian active learning with image data. International Conference Machine Learning (ICMI.), 2017.

---

Published as a conference paper at ICLR 2025
Emil Julius Gumbel. Statistical theory of extreme valuse and some practical applications. Nat. Bur Standards Appl. Math, Ser. 33, 1954.
Geoffrey E Hinton. Training products of experts by minimizing contrastive divergence. Neural computation, 14(8):177--180,, 2002.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Neural Infor- mation Processing Systems (NeurPPS), 2020.
Zhang-We Hong, Idan Shenfeld, Tsun-Hsuan Wang, Yung-Sung Chuang, Aldo Pareja, James Glass, Akash Srivastava, and Pulkit Agrawal. Curiosity-driven red-teaming for large language models. International Conference Learning Representations (ICLR), 2024.
Edward J. Hu, Nikolay Malkin, Moksh Jain, Katie Everett, Alexandros Graikos, and Yoshua Bengio. GFlowNet-EM for learning compositional latent variable models. International Conference on Machine Learning (ICML), 2023.
Edward J. Hu, Moksh Jain, Eric Elmoznino, Younesse Kaddar, Guillaume Lajoie, Yoshua Bengio, and Nikolay Malkin. Amortizing intractable inference in large language models. International Conference on Learning Representations (ICLR), 2024.
Moksh Jain, Emmanuel Bengio, Alex Grrnandez-Garche Jarrid Rector-Brooks Bonaventure FP
Dossou, Chanakya Ajit Ekbote, Jie Fu, Tianyu Zhang, Michael Kilgour, Dinghuai Zhang, et al. Bi- ological sequence design with gflownets. International Conference on Machine Learning (ICML), 2022.
Moksh Jain, Tristan Deleu, Jason Hartford, Cheng-Hao Liu, Alex Harnandez-GGrcia,, and Yoshua Bengio. Gflownets for ai-driven scientific discovery. Digital Discovery, 2(3):557-577, 2023.
Hyosoon Jang, Yunhui Jang, Minsu Kim, Jinkyoo Park, and Sungsoo Ahn. Pessimistic backward policy for GFlowNets. Neural Information Processing Systems (NeurIPS). 2024a.
Hyosoon Jang, Minsu Kim, and Sungsoo Ahn. Learning energy decompositions for partial inference of GFlowNets. International Conference on Learning Representations (ICLR), 2024b.
Hyeonah Kim, Minsu Kim, Sanghyeok Choi, and Jinkyoo Park. Genetic-guided GFlowNets for sam- ple efficient molecular optimization. Neural Information Processing Systems (NeurIPS), 2024a.
Minsu Kim, Sanghyeok Choi, Jiwoo Son, Hyeonah Kim, Jinkyoo Park, and Yoshua Ben- gio. Ant colony sampling with GFlowNets for combinatorial optimization. arXiv preprint arXiv:2403.07041,
2024bb. 2024b.
Minsu Kim, Joohwan Ko. Dinghuai Zhang, Ling Pan, Taeyoung Yun, Woochang Kim, Jinkyoo Park, and Yoshua Bengio. Learning to scale logits for temperature-condiiional GFlowNets. Interna- tional Conference on Machine Learning (ICML), 2024c.
Minsu Kim, Taeyoung Yun, Emmanuel Bengio, Dinghuai Zhang, Yoshua Bengio. Sungsoo Ahn, and Jinkyoo Park. Local search GFlowNets. International Conference Learning Representations (ICLR), 2024d.
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. International Conference on Learning Representations (ICLR), 2015.
Andreas Kirsch, Joost Van Amersfoort, and Yarin Gal. Batchbald: Efficient and diverse batch acquisition for deep bayesian active learning. Neural Information Processing Systems (NeurIPS), 2019.
Salem Lahlou, Tristan Deleu, Pablo Lemos, Dinghuai Zhang, Alexandra Volokhova, Alex Hemnndez- Léna Néhale Ezzine, Yoshua Bengio, and Nikolay Malkin. A theory of con- tinuous generative flow networks. International Conference on Machine Learning (ICML), 2023.
Elaine Lau, Nikhil Vemgal, Doina Precup, and Emmanuel Bengio. DGFN: Double generative flow networks. arXiv preprint arXiv:2310. /9685, 2023.
Elaine Lau, Stephen Zhewen Lu, Ling Pan, Doina Precup, and Emmanuel Bengio. QGFN: Control- lable greediness with action values. Neural Information Processing Systems (NeurIPS), 2024.
12

---

Published as a conference paper at ICLR 2025
Yann LeCun, Sumit Chopra, Raia Hadsell, M Ranzato, Fujie Huang, al. A tutorial on energy- based learning. Predicting structured data, 10), 2006.
Seanie Lee, Minsu Kim, Lynn Cherif, David Dobre, Juho Lee, Sung Ju Hwang, Kenji Kawaguchi, Gauthier Gidel, Yoshua Bengio, Nikolay Malkin, et al. Learning diverse attacks large language models for robust red-teaming and safety tuning. arXiv preprint arXiv:240S. 18540, 2024.
Pablo Lemos, Nikolay Malkin, Will Handley, Yoshua Bengio, Yashar Hezaveh, and Laurence Perreault-Levasseur Improving gradient-guided nested sampling for posterior inference. In- ternational Conference on Machine Learning (ICML), 2024.
Kanika Madan, Jarrid Rector-Brooks Maksym Korablyov, Emmanuel Bengio, Moksh Jain, Andrei Nica, Tom Bosc, Yoshua Bengio, and Nikolay Malkin. Learning GFlowNets from partial episodes for improved convergence and stability. International Conference on Machine Learning (ICML), 2023.
Nikolay Malkin, Moksh Jain, Emmanuel Bengio, Chen Sun, and Yoshua Bengio. Trajectory balance: Improved credit assignment in GFlowNets. Neural Information Processing Systems (NeurIPS), 2022.
Nikolay Malkin, Salem Lahlou, Tristan Deleu, Xu Ji, Edward Hu, Katie Everett, Dinghuai Zhang, and Yoshua Bengio. GFlowNets and variational inference. International Conference on Learning Representations (ICLR), 2023.
Charles C Margossian and David M Blei. Amortized variational inference: When and why? Uncer- tainty in Artificial Intelligence (UAT), 2024.
Volodymyr Mnih. Playing atari with deep reinforcement learning. arXiv preprint arXiv:1312.5602, 2013.
Sanmit Narvekar, Bei Peng. Matteo Leonetti, Jivko Sinapov, Matthew E Taylor, and Peter Stone. Curriculum learning for reinforcement learning domains: A framework and survey, Journal of Machine Learning Research, 21(181):1-55 2020.
Mizu Nishikawa- Toomey, Tristan Deleu, Jithendaraa Subramanian, Yoshua Bengio, and Laurent Charlin. Bayesian learning of causal structure and mechanisms with GFlowNets and variational Bayes. arXiv preprint arXiv:221 1,0763, 2022.
Frank Noé. Simon Olsson, Jonas Köhler, and Hao Wu. Boltzmann generators: Sampling equilibrium states of many-body systems with deep learning. Science, 365(6457) 1147, 2019.
Ling Pan, Nikolay Malkin, Dinghuai Zhang, and Yoshua Bengio. Better training of GFlowNets with local credit and incomplete trajectories. International Conference on Machine Learning (ICML), 2023a.
Ling Pan, Dinghuai Zhang, Aaron Courville, Longbo Huang, and Yoshua Bengio. Generative aug- mented flow networks. International Conference on Learning Representations (ICLR), 2023b.
Dominic Phillips and Flaviu Cipcigan. MetaGFN: Exploring distant modes with adapted metady- namics for continuous GFlowNets. arXiv preprint arXiv:2408, 15905, 2024.
Tom Schaul. Prioritized experience replay. International Conference on Learning Representations (ICLR), 2016.
Marein Sendera, Minsu Kim, Sarthak Mittal, Pablo Lemos, Luca Scimeca, Jarrid Alexandre Adam, Yoshua Bengio, and Nikolay Malkin. Improved off- off-policy training of diffusion samplers. Neural Information Processing Systems (NeurIPS), 2024.
Ozan Sener and Silvio Savarese. Active learning for convolutional neural networks: A core-set approach. International Conference Learning Representations (ICLR), 2018.
Kiyoung Seong, Seonghyun Park, Seonghwan Kim, Woo Youn Kim, and Sungsoo Ahn. Col- lective variable free transition path sampling with generative flow network. arXiv preprint arXiv:2405. 19961, 2024.
13

---

Published as a conference paper at ICLR 2025
Max W Shen, Emmanuel Bengio, Ehsan Hajiramezanali, Andreas Loukas, Kyunghyun Cho, and Tommaso Biancalani. Towards understanding and improving GFlowNet training. International Conference on Machine Learning (ICML), 2023.
Sam Sinai, Richard Wang, Alexander Whatley, Stewart Slocum, Elina Locane, and Eric D Kelsic. AdaLead: A simple and robust adaptive greedy search algorithm for sequence design. arXiv preprint arXiv:2010.02144, 2020.
Yang Song. Conor Durkan, lain Murray, and Stefano Ermon. Maximum likelihood training of score- based diffusion models. Neural Information Processing Systems (NeurIPS). 2021.
Zitao Song, Chao Yang. Chaojic Wang. Bo An, and Shuang Li. Latent logic tree extraction for event sequence explanation from llss. International Conference Machine Learning (ICML., 2024.
Daniil Tiapkin, Nikita Morozov, Alexey Naumov, and Dmitry Vetrov. Generative flow networks as entropy-regularized RL. Artificial Intelligence and Statistics (AISTATS), 2024.
Brandon Trabucco, Xinyang Geng, Aviral Kumar, and Sergey Levine. Design-bench. Benchmarks for data-driven offline model-based optimization. International Conference Machine Learning (ICML), 2022.
Austin Tripp. Erik Daxberger, and José Miguel Homández-Lobato. Sample-efficient optimization in the latent space of deep generative models via weighted retraining. Neural Information Processing Systems (NeurIPS), 2020.
Belinda Tzen and Maxim Raginsky. Neural stochastic differential equations: Deep latent Gaussian
models in the diffusion limit. arXiv preprini 1.XXvv:905..0 2019.
Nikhil Vemgal, Elaine Lau, and Doina Precup. An empirical study of the effectiveness of using replay buffer on mode discovery GFlowNets. arXiv preprint arXiv:2307.07777 2023.
Siddarth Venkatraman, Moksh Jain, Luca Scimeca, Minsu Kim, Marcin Sendera, Mohsin Hasan, Luke Rowe, Sarthak Mittal, Pablo Lemos, Emmanuel Bengio, et al. Amortizing intractable in- ference diffusion models for vision, language, and control. Neural Information Processing Systems (NeurIPS), 2024.
Cheems Wang, Yiqin Lv, Yixiu Mao, Yun Qu Yi Xu, and Xiangyang Ji. Robust fast adaptation from adversarially explicit task distribution generation. arXiv preprint wxXiv22407.19522 2024.
Donggeun Yoo and In So Kweon. Learning loss for active learning. Computer Vision and Patlern Recognition (CVPR), 2019.
Fangxu Yu, Lai Jiang, Haoqiang Kang, Shibo Hao, and Lianhui Qin. Flow of reasoning: Training Ilms for divergent problem solving with minimal examples. arXiv preprint arXiv:2406.05675 2024.
Dinghuai Zhang, Nikolay Malkin, Zhen Liu, Alexandra Volokhova, Aaron Courville, and Yoshua Bengio. Generative flow networks for discrete probabilistic modeling. International Conference on Machine Learning (ICML), 2022.
Dinghuai Zhang, Ricky Tian Qi Chen, Cheng-Hao Liu, Aaron Courville, and Yoshua Bengio. Diffu- sion generative flow samplers: Improving learning signals through partial trajectory optimization. International Conference Learning Representations (ICIR), 2024.
Qinsheng Zhang and Yongxin Chen. Path integral sampler: stochastic control approach for sam- pling. International Conference on Learning Representations (ICLR), 2022.
Shuo Zhang, Yang Liu, and Lei Xie. Molecular mechanics-driven graph neural network with multi- plex graph for molecular structures. arXiv preprint arXiv:2011.07455 2020.
14

---

Published as a conference paper at ICLR 2025
A LIMITATIONS
The primary limitation of the adaptive Teacher the added complexity in training, as the Teacher's policy network must be trained in addition to the Student's. This introduces trade-off' between increased training complexity and enhanced mode-secking capabilities. While we believe that the ability to discover multiple modes outweighs the additional complexity, it is important to apply this technique judiciously. For tasks where the reward is unimodal or that do require extensive exploration, using a Teacher may not be necessary. However, in scenarios where the model tends to collapse to specific modes of the multimodal target distribution, employing Teacher is beneficial choice.
Additionally, since the Student struggles to cover entire modes on its own, in significantly larger search spaces the Teacher may also have difficulty covering the full range of modes that the Student fails to discover, potentially collapsing into specific modes. In such large-scale settings, we expect that multi agent Teacher system with multiple agents collaboratively covering the space- could be a beneficial direction for future work to mitigate this limitation.
B THEORETICAL ANALYSIS OF STATIONARY DISTRIBUTIONS
Proposition 1. Let the behavior policy PT(T) be distribution over trajectories E T that satisfies
full support.
the parameters 8" and o" of the Student and Teacher policies, respectively, jointly optimize the objective functions to in expectation over P3(T), then:
(a) The marginal distribution of the Student policy over terminal states satisfies
Oxxxx0") & R(x),
(b) The marginal distribution of the Teacher policy over terminal states satisfies
(xxxx'') oc R(x)",
where:
R(x) is the reward function, functionn
is offset constant introduced in (5),
a > 0 is the mixing constant for the Teacher introduced in (6).
Proof: According to the trajectory balance training theorem (Malkin et al., 2022), the Student policy for all r) and only if its marginal distribution over achieves optimal loss (i.e., LTB(T; terminal states pxx, 0*) x R(x).
Suppose now that the Student policy has reached this optimum. Then the reward for the Teacher is
log Teacher (x; 0²) = (P(x()) log (e (1+ C Is(r. 0(): a log R(x)
= -Bpa(r)x:0) [log(e)) a log R(x)
-lo(()) log R(x)
= =log(e((x)").
(10)
Again by the TB training theorem, the loss is minimized if and only if the marginal distribution of the Student policy over terminal states proportional to e <R.- equivalently, to R(.)a
15

---

Published as a conference paper at ICLR 2025
C DETAILED IMPLEMENTATION
For all tasks, we set a = 0.5 except for the exploration iintensinn deceptive grid world tasks, where 0.0. C is set to 19 for all tasks. Our hyperparameter analysis is provided in Ap- pendix E.4, In a diffusion sampling task, to concentrate the teacher on high-reward regions within a vast continuous space, we limit the teacher's training set to samples where x > "threshold- Here, Fihreshold the 90th percentile reward from the untrained student's samples.
For neural networks, we use identical architectures for both the Student and Teacher models. Specif- ically, for the GFN architecture design, we match the architectures used by each baseline for every task. Detailed descriptions are provided in Appendix D.
Regarding the replay buffer, we adhere to existing implementations for each task and follow pri- oritized replay rules. In grid world and diffusion sampling tasks, we use rank-based priority as introduced by Tripp et al. (2020) and Sendera et al. (2024). For biochemical tasks, we employ portion-wise priority as proposed by Shen et al. (2023). Additionally, we implemented two variants of loss-prioritized buffer: one prioritizes based on TB loss, and the other uses Teacher as the prior- ity, serving similar purpose. We analyze the performance differences between these variants and refer the use of Reacher as "PER."
When selecting the behavior policy during training, we periodically choose among the Student, Teacher, and buffer in specific proportions: ratio of 1:1:0 for Grid World tasks, 3:1:2 for Diffusion Sampler tasks, and 2:1:3 for Biochemical tasks. In the diffusion sampler, we use PER for the replay buffer. For the biochemical task, we benchmark the teacher using both PER (referred to as PER + Teacher) and reward- priioriidpi buffer (referred to PRT + Teacher).
higher Student ratio encourages exploitation, higher Teacher ratio promotes exploration, and a larger buffer enhances sample efficiency until reaching the replay ratio barrier (D'Oro et al,, 2023).
In deceptive grid world tasks, exploration is crucial, so we assign a higher Teacher ratio. For dif- fusion sampling tasks, using a prioritized replay buffer with Teacher rewards yields strong perfor- mance. Blending this approach slightly with the Teacher enhances both performance and conver- gence speed, achieving mode coverage for both EUBO and ELBO. biochemical tasks, sample efficiency is critical. Therefore, existing tasks are set to be on-policy with buffer ratio of 1:1. We adjust the proportions to favor on-policy learning by setting the crudent-to-teacht ratio 2:1, as biochemical tasks require some level of exploitation, as reported by Shen et al. (2023).
Other implementation details and hyperparameters are maintained as in each task's ' experimental setting.
D DETAILED EXPERIMENTAL SETTING
D.I DECEPTIVE GRID WORLD
Evaluation metrics. Following Bengio et al. (2021),
Table 3: The total number of terminal states ((XI) and modes ([M)) for cach grid setting, where d is the dimension and H is the hori- zon the hypergrid.
use the number of modes discovered and the em-
pirical L1 distance from target distribution as evalua- tion metrics. We define the modes x's with R(x) R₂ + Rg. The exact size of the search space |X| is (H 11-11 1)d d(H - 1), with d and I/ representing the di- mension and horizon of the hypergrid, respectively. [X| and the total number of modes [M| for each grid setting are reported in Table 3. The L.1 distance calculated by [XT I Exex [xp=(x) R(x)/Z]], where Z Σ, R(x), which known in this synthetic task. Unlike previous works (Bengio et al., 2021; Malkin et al., 2022) where portion of the final training samples was used to approximate the expectation, we generate 10s new
X
IM|
d = 2, 11 128
16383
676
= 2, 11 256
65535
2601
d 11 16 d 4=
64125
81
1042685
1296
samples from policy calculate (x(x) for evaluation. We use one sample for only one gradient step by default, i.e., update-to-date (UTD) ratio (Chen et al., 2021) is 1, except for the case using replay buffer, where the ratio increases to 2.
Hyperparameters. Following Bengio et al. (2021); Malkin et al. (2022), we use two-layer MLP with 256 hidden units for the parameterized policy Pp (::0) along with a learnable parameter for log Zg. We train them using the Adam optimizer with learning rate of 10-3 for policy and 10 for log Zg. The backward policy PB is fixed as a uniform random policy. When using replay buffer,
16

---

Published as a conference paper at ICLR 2025
its size is dynamically set to [0.1]X]]. The total reward call budget is capped at 96,000, except for (d = 4, 32), which is increased to 384,000, considering the significantly larger search space. We use batch size of 16. The total number of gradient steps equals the number of reward calls divided by the batch size, and this number doubles when the replay buffer used.
Baseline implementations. In our experiments, we set for e-exploration to 0.01, as this value yielded the lowest L1 eπror among the tested options {0.001, 0.003, 0.01, 0.03, 0.1) in the setting where d and 128. Similarly, for GAFN (Pan al., 2023b), we set the intrinsic reward seale to 0.01, which also resulted in the lowest L1 error among the values (1.0, 0. 0.0, 0.001) under the same conditions. For PRT and PER, we use the same buffer size and prioritization scheme as above. Note that, unlike the original PER for malue-based RI., the buffer we used contains only terminal x rather than all state transitions. From a sampled x, the trajectory can be constructed by backward generation using PB. We also explore the integration of transition- replay buffer with the detailed balance objective in Appendix F.
D.2 DIFFUSION SAMPLING
In our diffusion sampler, primarily adhere to the settings of Sendera et al. (2024), which in turn build upon those of Zhang & Chen (2022).
Tasks. We benchmark the following two tasks:
Gaussian Mixture Model with 25 Modes (25GMM). The 25GMM consists of a two-dimensional Gaussian mixture with 25 modes, each having variance of 0.3. The mode centers are positioned on the grid (-10,- -, 0, 5, 10) x {-10, 0, 10).
Manywell (Noé er al., 2019). The Manywell is a 32-dimensional distribution formed as the prod- uet of 16 identical two-dimensional double-well distributions. Each component is defined by the potential function
(xx, x2) exp (-x 6x 0.5xx 1..5x)-0.500 -0.5x3)
(11)
Evaluation metrics. We measure evidence lower bound (ELBO), importance sampled ELBO (ELBO-IS), and evidence upper bound (EUBO).
For estimating ELBO, we draw samples from current policy and take average value of estimated log Z, which is log R(.) log (.(.) log PF(+) as follows:
M Σ 1=1
ELBO
(log R(x)) log Pb(r log PF(T':0), +- Pp::)), - (12)
M
where r/ mo x means that is the final state of r!.
Calculation ELBO-IS is similar ELBO:
M
ELBO-IS z log M Σ
exp(log R(x)) + log (n(T((x)) - log P)( +'~ PP(;;))) m xf-
(13)
EUBO was introduced as a metric that measures mode coverage by Blessing et al. (2024). To calculate EUBO, sample M samples from the target distribution and take the average of their variational log-likelihood bounds as follows:
M EUBO = M Σ (log R(x)) + log Pn(T'(x) log Pr(T':0) Pf))'((x), ~ PA(T xf). (14)
Forward and backward transition modeling. The diffusion sampler models discretized SDE trajectories T = (Xo + XAr + .... + x1), starting from Xn = (0, 0). Here, At = 1/T, where is the number of discrete time steps.
The forward policy Pp(x++, x;; 0) is modeled Gaussian distribution with mean x, uxx, 1; A) and covariance "
Pp(x xx:0) = N (Kraari x, u(xx,,;))) dPArT).
(15)
17

---

Published as a conference paper at ICLR 2025
Here, ux(x, r, 0) is the learnable score function, o is the standard deviation, and I denotes the identity matrix to ensure isotropic covariance.
The backward policy (B( Arr | xx) is defined as a discretized Brownian bridge with noise rate d:
PB(x xx) = Xy-ari Xx. I- Ar 1-A0o A'Ar1).
(16)
The densities of the distributions over complete forward and backward trajectories are given by:
T-1
Pr(T: 0) = [
Pr (F((())) 0). Pa(r x1)
[ i=1
Pn(xire (x(+t)ar)
(17)
i=0
The diffusion policy parameter 0 is trained using the TB loss.
Hyperparameters. We set x² 5.0 for 25GMM and or 1.0 for Manywell, with the number of time steps T = 100, following Sendera et (2024). We employ the same architecture as Zhang & Chen (2022) and Sendera et al. (2024), increasing the hidden dimension from 64 to 256 for Manywell accommodate the 32-dimensional tasks, and apply this adjustment all baselines. For replay buffer capacity, we set it to 5,000 for the 25GMM task and 20,000 for the Manywell task. All learning hyperparameters remain identical those Sendera et al. (2024). For evaluation, we set the number of samples to M = 2, 000.
D.3 BIOLOGICAL AND CHEMICAL DISCOVERY
For biochemical design tasks, we mostly follow the setting of Shen et al. (2023). For tasks, we use mrepend-appent MDP (PA-MDP), where the action is defined as adding token at the beginning or the end of a partial sequence, This MDP formulation makes possible to have multiple trajectories associated with the same design configuration
Hyperparameters. For training GFlowNets, we use similar setting proposed by prior works (Shen et al., 2023; Kim et al., 2024d). We use Adam optimizer (Kingma Ba, 2015) with learning rate 10-2 for log Zo. 10- for forward policy, 5 x 10 for teacher policy. To parametrize forward policy, we adopt relative edge flow policy parametrization mapping (SSR) from Shen et al. (2023). For QM9 and sEH tasks, we employ two-layer architecture with 1024 hidden units, while for the other tasks, we choose use two-layer architecture with 128 hidden units. We initialize log Zo to 5.0 for all methods. For backward policy, we use a fixed uniform policy. In terms of reward exponent, use value of 20 for both QM9 and TFbind8. For sEH and L14-RNAI, use relatively higher values, 200 and 40, respectively.
Evaluation metrics. For evaluation, we compute the number of modes using all the samples col- lected over the course of training. What we count as mode should have high reward. Unlike with other designs, we define "mode" as a configuration whose reward is above certain threshold. This is different from previously used metrics of mode counting assess diversity. For QM9 and TFbind8, we use a default mode set suggested by (Shen et al., 2023). For sEH, we set the reward threshold the top 0.01% % of X in terms of the reward and the diversity threshold as 0.4 Tanimoto diversity. For L14-RNAI, we set the reward threshold as the top 0.01% % of X in terms of the reward and the diversity threshold are described in Appendix D.2. We generate M 2, 048 samples from both the trained policy and the target distribution for the ELBO and EUBO calculation. To sample from the target distribu- tion, we evaluate all possible configurations and sample them proportionally to their reward. As the number of possible configurations is enormously large for sEH and L14-RNAI tasks, we use the Gumbel -mx trick (Gumbel, 1954) for sampling from the discrete probability distribution.
unit of Levenstein distance. We also report ELBO and EUBO which
18

---

Published as a conference paper at ICLR 2025
Deceptive grid world (d 2,/1 = 128)
# #oodss discovered (1)
x10
600
distance (L)
1.00
400
200
0.50
Num. reward calls
x10
Num. reward calls
x10
Deceptive grid world (d 2, 256)
x10
# modess discovered (t)
3.0
2000
L1 distance ())
2.5
2.0
1000
1.5
0.0
0.2
0.4
0.5
0.8
1.0
0.0
0.2
0.4
0.6
0.3
10
Nom. reward calls
-1
Num. reward calls
x10
Deceptive grid world (d 4, 16)
#1 modes discovered (t)
1.20
x10
30
40
distance
1.10
1.00
Num. reward calls
Num. reward calls
x10
Deceptive grid world (d 4,H H 32)
x10
# modes discovered (1)
300
distance (L)
200
1.68
1.64
1.63
Num. reward calls
X0
Num. reward calls
X10
TB
TB GAFN
TH PRT
TB PER
1B Teacher
TB Tescher PER
Figure 6: Evolution of evaluation metries for each method in deceptive grid world task. The value with standard deviation is depicted from five independent runs.
E ADDITIONAL EXPERIMENTAL RESULTS
E. EXTENDED EXPERIMENTAL RESULTS DECEPTIVE GRID WORLD
Fig. 6 shows the changes in two evaluation metrics during training: the number of modes discovered and the empirical /-, distance between the target and sampled distributions. The results indicate that Teacher highly effective at discovering modes and also generally performs well in terms of the I. distance. In contrast, while PER and PRT improve mode discovery, they tend to slightly worsen L performance. We believe this is because the replay buffer provides regularizing signal that prevents mode collapse, making the policy learning process more challenging compared to using purely on- policy methods without regularization. This leads to slight underfitting. On the other hand, on on-policy methods can easily overfit to the target distribution (but into specific modes), resulting in descent Li performance. However, without off-policy regularization, they tend to drop modes and fail to cover all modes in the target distribution. Teacher achieves high performance in both L and mode coverage. indicating that it not only offers off off-policy regularization ensure comprehensive mode coverage but also provides efficient curricula for faster convergence to the target distribution, leading strong performance across both metrics.
19

---

Published as a conference paper at ICLR 2025
Table 4: The effect number Monte Carlo samples (NMC) deceptive grid world task.
Grid config.
d
256
d
32
Algorithm Metric -
# modes (1)
Lj x 10 (J)
modes (7)
L1 x 10 (()
Nuc I (default)
2452.6 a 21.7
0.94 0.03
246.6 + 14.7
1.634 0.001
Mvc =
2489.6 = 15.8
0.96 0.07
234.6 14.2
1,635 0,000
Nac =
2490.4 # 30.3
0.94 + 0.06
230.8 = 3,4
1,634 = 0.000
NMc 10
2492.0 +: 26.3
0.96 1 0.05
239.0 + 11.8
1.634 0.000
Table 5: Ablation study on C in deceptive grid world task.
Grid config.
d = 2, 256
d= 4, 32
Algorithm Metric
# modes (1)
Lj x 10 (L)
modes (1)
L X 10-5 ())
C=0
2469.4
1.06
253.2
1.634
C
2454.2 : 18.7
0.95 0.02
256.0
1.635 0,000
C 19 (default)
2452.6 21.7
0.94 + 0.03
246.6 14.7
1.634 0.001
C 29
2465.2 a: 30.6
0.94 0.04
243.4 8.5
1.634 0.000
E.2 STUDY ON MONTE CARLO APPROXIMATION FOR RTEACHER
We use Monte Carlo estimate with a single trajectory to approximate the expectation in Eq. (4) and Eq. (5). To validate that this single-sample approximation is reasonable, we test our algorithm in the deceptive grid world task with an increased number of samples, ranging from 10. We do not use the replay buffer in this analysis. As described in Appendix E.1, the performance is not significantly affected by NMC- This suggests that using Monte Carlo approximation with sample size of to estimate the stochastic reward Reeaher was reasonable.
E.3 ABLATION STUDY ON THE CHOICE OF C VALUE
We set the hyperparameter C in Eq. (5) to across all experiments without an extensive hyperpa- rameter search. To evaluate this choice, we perform an ablation study on the deceptive grid world task, testing alternative values of 0, 9, and 29. Note that the replay buffer not used for this analysis. The results are summarized in Table 5. Although the effect on the number of modes discovered is somewhat mixed,
performs slightly worse than the other values regarding the empirical error, supporting our hypothesis that focusing more on undersampled regions is beneficial. We also found that the results are not highly sensitive the choice of C.
20

---

Published as a conference paper at ICLR 2025
E.4 STUDY ON REWARD MIXING
We introduce a to mix the reward based on the Student's loss and Student's log reward help Teacher target both high-loss and high-reward areas effectively. In this section, we investigate the effect of a on the performance of our method.
Table 6: Mixing component study on deceptive grid world task
Grid config
2.
256
d 4, 32
Algorithm Metric
modes (1)
10
modes (?)
1,j 10 (L)
0.0
2452.6
0,94 =
246.6 14.7
1.634 0.001
0.5
2415.2 # 262.8 0,90
8.56
1.634 1.000
Deceptive grid world. Table 6 shows that reward mixing with = 0.5 degrades performances for mode seeking high dimensional tasks as deceptive grid world task is exploration intensive task; teacher solely focusing on high loss region is more beneficial. Still mixing with a = 0.5 outperforms other baselines.
Table 7:
Mixing component study on diffusion sampler task
Energy
25GMM (d 2, log Z 0)
Manywell (d 32, log Z 164.696)
Algorithm Metri
ELBO (7)
ELBO-IS (f)
EUBO (L)
ELBO
ELBO-IS (f)
EUBO (1)
a 0.0
-0. 10.4400000
~0.009
0. 0.122a0010
163 447
164, 694.000
166. 066.0242200
-0.137-0001
+0.005
0..15
163.484.00m
164.676-00x
165. 800-000
Diffusion sampling. As shown in Table 7, mixing with a = 0.5 shows slightly better performance, though both achieve significantly higher sampling efficiency compared to the baselines. Both 0.0 and = 0.5 come close to reaching the target log Z.
TFbind8 (EUBO(() (J))
TFbind8 (Num Modes (f))
300
11
100
200
400
600
800
1000
200
600
800
1000
Num. Training Rounds
Num. Training Rounds
TB Teacher PER (a=0.0)
TB Teacher PER (om0..5
Figure 7:
Training graph for TFbind8 task by varying a of reward mixing. Mcan value with standard deviation depicted five independent runs.
Biological and Chemical Discovery (TFbind8), As shown in Fig. 7, mixing with a = 0.5 yields significantly better performance in TFbind8 task. This highlights the importance of having the teacher focus on both high-loss and high-reward areas.
21

---

Published as a conference paper at ICLR 2025
# modes discovered (t)
L1 distance (L)
X10
1.600
1.650
400
1.645
200
1.635
1.630
0
1.255
Num. reward calls
X10
Num. reward calls
X10
TB + PER
TB- Teucher + PER
TB LS
PER +
TB Teacher + PER + LS
Figure 8: Evolution of evaluation metrics with or without the local search in deceptive grid worlds (d
32). The mean value with standard deviation depicted from five independent runs.
Table 8: Comparison with local search (LS) (Sendera et al., 2024) on Manywell
Energy -
Manywell (d = 32, log 7. = 164.696)
Algorithm 1
ELBO (T)
ELBO-IS (1)
EUBO (1)
PER
161.537=0.186
162.58220.266
210.440aa.Hss
PER + LS
163.308=0.189
164.508a0.66
168.953+3.209
Teacher
163.484-0049
164.67 04:004
165.800e0.045
Teacher + LS
163.472200066
164.685=00063
165.787#0.044
E.5 TEACHER WITH LOCAL SEARCH
Local search is useful technique to improve the sampling quality of GFlowNets (Hu et al., 2023; Kim et al., 2024d). In this section, we investigate the possible integration of local search and Teacher and compare it with existing local search integrated solely on Student.
Deceptive grid world. We tested the backtracking anderinnnaanddrrcoossttttiiooo local search method intro- duced in $3.3 in the deceptive grid world, using grid configurations of (d = 2, Il = 256) and 4d 32). Four iterative local searches were performed every 16th training batch. The back- tracking ratio is 0.5, meaning the last half of a trajectory is destroyed and reconstructed by policy. We applied a deterministic acceptance rule, accepting a new trajectory if it had a higher Reacher- For comparison, we used two baselines: on- on-policy TB and TB with PER, both using the same local search but using the task reward R to determine the acceptance.
The experimental results extending is illustrated in Fig. 8. Teacher with PER outperforms both PER and on-policy TB with a large margin terms of mode coverage, regardless of whether local search is applied. When combined with local search, Teacher achieves the best results in both the number of modes discovered and the empirical L1 distance. We believe this performance gain is largely due to the reduction of non-stationarity $3.3, though isolating the exact contribution is complex and left for future work.
Diffusion sampling. We utilize the Manywell task to compare the effect of local search on the Teacher model. Specifically, we employ parallel local search methods (Sendera et al., 2024) that leverage eetropoois-Haasingggggieee Langevin dynamics (MALA) on samples from the replay buffer to refine sample quality. As shown in the Table 8, by integrating this local search with Priori- tized Experience Replay (PER), we observe significant performance improvements.
Remarkably, the Teacher model even without local search still outperforms these results. This is notable because the Teacher's exploration does not require gradient information of the energy function, whereas MALA relies such gradient information. Since the Teacher rapidly achieves optimal sampling quality on the Manywell task, we observe not much improvement when applying local search to the Teacher in the diffusion sampling task.
22

---

Published as a conference paper at ICLR 2025
TFaind8 (EUBO (J)
TFbind8 (Num Modes (1))
300
100
200
400
600
1000
200
400
800
1000
Num. Training Rounds
Num. Training Rounds
TB + PER
TB. Tencher + PER
-----
TB PER LS
TB Teacher PER LS
Figure 9: Training graph for TFbind8 task by integrating local search (Kim et al., 2024d). Mean value with standard deviation is depicted five independent runs.
Biological and Chemical Discovery. We utilize the TFbind8 task to compare the effect of local search the Teacher model. Specifically, employ the local search method suggested by (Kim et al., 2024d), which involves backtracking and reconstructing sequences using the forward and backward policies of GFlowNets. The decision accept adjusted samples is based on whether R(x') R(x), where is the new sample. For the Teacher model, as described in $3.3, we perform local search to mitigate non-stationary in the student and to optimize Reeacher For both the Student and Teacher models, we employ Prioritized Experience Replay (PER).
We visualize the results Fig. 9. As shown in the figure, the Teacher model without local search still outperforms the Student model with local search. This demonstrates that the exploration capa- bility of the Teacher model is far more efficient than conducting local search with the current policy. Moreover, we observe that integrating local search with the Teacher model leads to further improve- ment in terms of EUBO, highlighting the synergistic effect of combining Teacher model with local search.
23

---

Published as a conference paper at ICLR 2025
Student (0/5)
Student (1/5)
Student (2/5)
Student (3/5)
Student (4/5)
Student (5/5)
Teacher (0/5)
Teacher (1/5)
Teacher (2/5)
Teacher (3/5)
Teacher (4/5)
Teacher (5/5)
Figure 10: Illustration of the distribution dynamics between the Teacher and Student models, along with their stationary distributions. The Student (ratio) represents the fraction completed training steps.
Student (0/5)
Student (1/5)
Student (2/5)
Student (3/5)
Student (4/5)
Student (5/5)
Teacher (0/5)
Teacher (1/5)
Teacher (2/5)
Teacher (3/5)
Teacher (4/5)
Teacher (5/5)
Figure 11: Illustration the distribution dynamics between the Teacher Student models, along with their stationary distributions. The Student (ratio) represents the fraction of completed training steps.
E.6 2D PLOTS OF TEACHER AND STUDENT OVER TRAINING
This section presents the distributions of the Teacher and Student models during training. For this visualization, set the mixing component explore the Student's loss regions. As shown in Fig. 10 and Fig. 11, the Teacher effectively identi- fies the Student's 's missing modes, providing suitable training distribution throughout the epochs, ultimately enabling the Student to discover all modes successfully.
O, meaning the Teacher's major objective is to
24

---

Published as a conference paper at ICLR 2025
Table 9: Evaluation results of DB algorithms deceptive grid worlds with dimension d and grid length H. Mean and standard deviation from independent rums are reported. bold applied to the best mean value among DB-based methods.
Grid contig.
d 2. II 128
d 2, H 256
d 4, 16
d 4, 32
Algorillum Metric
modes
x0-10
kuudles
L 10-2
#oodes
1
#modes
10
TB
(on-policy -)
645 4
0.58
733.6
6.6
1,027 2.012
16.6
1.635 a
Teacher
676.0
2.13
2452.6221.7
a.94
41.4
1.19 0.16
24.
1.634 0001
DB
(on-policy
644.0 13.3
=.57
1025.4
1 0.0
1.7
0.009
6.8
1.634 0:00
e-expl.
578.2
2.4
1.10
0.004
25.2
1.635
PER*
316.6 166.5
6.01 =062
890.4 241.7
1,75 0.00
2.6
1,005 0.012
17.2
1.634 0.00
Teacher
675.4
4.42
1817.6 49.8
1.59 0:00
$8.2
1,009 = 0.007
1.632 0,00
F TEACHER FOR DETAILED BALANCE
In this section, we extend the proposed idea in $3 to detailed balance (DB; Bengio et al., 2023), another GFlowNet learning objective.
F.1 DETAILED BALANCE AND TEACHER'S 'S REWARD FOR DB
The general problem settings, including MDP formulation and reward function, are the same as $2. Unlike TB, which requires parameterizing Zo. DB parameterizes the state flow function F(s;0) for each state, along with PF and PB- Note that F(x; () R(x) for every terminal state x e X, and the initial state flow F(so) is an estimate of the total reward. The detailed balance discrepancy defined for any transition of states (s, a,,'s)) as
ōpa a, s' = [log F(s': 0) + log (n(s) (s:0)] - [log F(s; log Pr(s' sx; 0)]
(18)
backward edge flow
forward edge flow
Same TB, when öpa a, :) for all (s,a,s'), then Px(x) R(x) is achieved for all x, where PT defined as Eq. (1). We can naturally define DB loss as ōDB (s, a, s'; 0)2 on each transition (s, s') sampled from behavior policy л. For more formal derivation, please refer to Bengio et al. (2023).
Analogous to Eq. (4) and Eq. (5), we define the basic form as and the re-weighted version of Teacher's 's reward for DB. The basic form is
log R basic
Teacher-DB (x; 0) = Er((T(s)
log Σ (S.
ÔDB (s.a, (s.a, s';
(19)
and the re-weighted version is
log
weighted Teacher-DB (x;6)= = =Epo(T)x:0)
log E+
Σ (s,0,5))
(1+Claw(ssa ōDB (s, a, 0)3
(20)
where we approximate the expectation using single trajectory.
The reward mixing in Eq. (6) is not directly available in the DB case since it entails a non-trivial credit assignment problem. Thus, we set RTeacher DB = weiggtted "eacher--B
The overall training procedure is similar to $3.2 and Algorithm 1, except we use DB loss for both Student and Teacher training. For a given transition (s, a, ~ π, the DB-loss functions are defined by
Lstudent-D (0, s': 6DB a, -log ((s 1s:
(21)
T(s':0)Pn(s' 15)
LTeacher-DB as DD s'; |og F(s: s;
(22)
F(s': Pn(s (s)
where F(x; 0) R(x) and F(x; Φ) Reeacher DB-
25

---

Published as a conference paper at ICLR 2025
F.2 EXPERIMENTS IN DECEPTIVE GRID WORLD
We use the same experimental settings as $5.1. We incorporate a transition-based replay buffer, meaning that we save all state transitions along the trajectory rather than saving only the terminal state x as in the TB case, This allows closer implementation of the original PER, where the prioritization is performed with TD error for each transition. To distinguish the PER used in $5.1, we call the transition-basse based PER as PER* Regarding the baselines, we do not benchmark PRT is not trivial to assign an episodic reward at the terminal state to each transition. We omit the GAFN since its source code only supports the TB algorithm. We also include TB, TB with Teacher for reference.
The result is summarized in Table 9. Similar to the TB case (Table 1), Teacher provides significant improvement over baselines for DB. This confirms that our method offers flexibility across different GelowNets objective functions.
G SCALING EXPERIMENTS
In this section, we demonstrate the scalability of our method. We first test it on larger-scale tasks on Deceptive Gridworlds, show that its effectiveness remains consistent as the scale increases. Then we apply our method to a real-world task of prompt sampling on large language models (LLMss, where we discover desirable prompt sentences that require effective exploration of the combinatori- ally large language search space.
Table 10: Evaluation results large- scale deceptive grid worlds with dimension d and grid length H. The mcan and standard deviation of the number of modes discovered (# modes) from 3 independent runs are reported. Due to the computational expense obtaining the exact target distribution in large-scale problems, L distance is excluded from the analysis. The best mean values are highlighted in bold, while the second-best are marked with an underline.
Grid config.
d d= = 64
d= 4, 11 128
d, 11 32
d 6, 11 64
Num. terminal states |X| -
1.68 x 107
2.68 x 10
1.06
10
6.85 10
TB
(on-policy -)
2.7
228.7 38.1
0.3 0.5
3.7 2.6
e-expl.
49.7 17.2
866.7 154.4
0.7
11.3 4.0
GAFN
42.0
180.0 51.9
1.3
4.0
PRT
119.3 + 16.0
222.7 + 11.6
6.7
0.9
PER
70.7 + 12.3
164.3
2.0
0.9
Teacher (ours)
299.0 + 10.7
728.0 # 192.0
9.7 3.4
21.3 + 6.0
G.1 LARGE-SCALE EXPERIMENT ON DECEPTIVE HYPERGRIDS
We evaluate our algorithm larger-scale settings of deceptive grid worlds. Details of the grid configurations and the total number of terminal states (representing the size of the search space) are provided in Table 10. We use the same experimental settings used for the grid with (d = 4, H 32), which are described in Appendix D.I. Note that we omit L1 distance from our metrics, as calculating the exact target distribution is computationally infeasible for these large-scale problems.
As shown in the results Table 10, our algorithm generally outperforms other baselines even in envi- ronments with a much larger search space. We conjecture the strong performance of e-exploration in the (d 4, H 128) configuration is because, when H is large, clusters of adjacent high high-reward states are large. In such cases, the random, brittle actions from e-exploration can be advantagcous in discovering multiple adjacent modes within the same region.
G.2 SAMPLING LLM ATTACK PROMPTS
We benchmark our off-policy training method on the automated red-teaming task using GFlowNets, following the approach of Lee et al. (2024). In this task, the problem is formulated as inference over prompt sequences proportional to a target reward, which is computed by a toxicity classifier evaluated on the response the prompt induces from fixed vietim model. We adhere the same settings, baselines, evaluation model, and model architecture as the mentioned work. Specifically. fine-tune GPT-2 GFlowNet policy serve the attack model.
Setting. The log-reward is defined weighted mixture of the log-likelihood of the language model and the toxicity score of the prompt. Toxicity evaluated by classifier that measures how likely the prompt is induce toxic outputs from the victim model.
26

---

Published as a conference paper at ICLR 2025
1.0
0.92
0.91
Success Rate
0.90
0.6
GFN attacker
0.59
(ours)
0.58
REINFORCE
PPO Novelty
0.87
9.86
GFN attacker
SFT
attacker Teacher
0.0
0.85
0.815
0.630
0.635
0.640
0.645
0.500
Diversity
Diversity
(a) vs. All baselines
(b) vs. GFN methods
Figure 12: Percentage of toxic prompts (y-axis) versus prompt diversity (x-axis, measured 1 1-cosine similar- ity) for attack methods on GPT-2.
Baselines. We compare our method with several baselines: In-Context Learning (ICL), Supervised Fine-Tuning (SFT), REINFORCE, and Proximal Policy Optimization (PPO) with novelty reward as suggested by Hong et al. (2024). We also include the GFlowNet attacker proposed by Lee et al. (2024), which the waae-offthe-aad method leveraging GFlowNet sampler with sophisticated off- policy techniques using replay buffer.
Implementation. We implement a teacher network over the GelowNet attacker. While the origi- GelowNet attacker uses on-policy updates and 1:1 ratio in the replay buffer, we use Teacher, Student (on-policy), and replay buffer trajectories in a 2:1:3 ratio. We aim to observe whether this adjustment provides any benefits over the baseline. We reproduce the GFN attacker results with ß E 0.06, 0.07, 0.08, where ß is the temperature parameter for the toxicity reward. For our teacher method, we test with ß e 0.02, 0.05. For other baseline results, we directly use the data reported in the figures by Lee et al. (2024).
Results. As shown in Fig. 12a, the teacher network achieves slightly higher diversity and success rates compared the state-of-the-art GelowNet attacker. Other baselines fail produce both di- verse and toxic prompts: REINFORCE leads to mode collapse, and ICL and SFT do not generate meaningful toxic prompts (see Lee et al. (2024) for more detailed analysis of these baselines). The GelowNet attacker provides well-balanced results, achieving high toxicity in successful prompt sen- tences with diversity. Our Teacher network offers slight improvement attacker method (see Fig. 12b), demonstrating that our approach can be flexibly applied to real-world tasks. Notably, even when using a lower & (indicating a peaky toxicity reward) than the GFlowNet attacker. the diversity achieved is higher. This suggests that the teacher encourages exploration into missing modes to enhance diversity.
the basleine GFlowNet
These findings show the potential applicability of the teacher concept in large language model rea- soning tasks where amortized inference using off-policy RL has been applied, including automated red-teaming, infilling, chain-of-though reasoning, and planning (as studied in Hu et al., 2024; Song et al., 2024; Yu et al., 2024).
27


Reviews:
[{"review_id": "F0NRIk6bBd", "reviewer": "Reviewer_tjcQ", "summary": "This paper proposed a method to improve mode coverage and training efficiency in amortized inference methods like GFlowNets. Specifically, the authors use offline RL training to encourage the discovery of diverse, high-reward candidates, and addressed the key challenge -- exploration in off-policy RL. The main idea is to use an adaptive \"Teacher\" model to help the \"Student\" sampler by focusing on regions with high loss. The teacher model is used as as an auxiliary model, is trained to target areas where the Student model has high errors. This allows it to cover unexplored modes (which usually have high errors) and provide a more efficient training process. Empirically, the authors show that this approach works well in various tasks, such as discrete sequence design and continuous diffusion sampling, with better sample efficiency and mode coverage.", "questions": "Can the author explain why you chose GFlowNets for the experiments? Are they more effective than diffusion models for chemical or drug discovery? In my understanding, diffusion models can fairly easily to fit multimodal distributions.", "limitations": "", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The idea of exploring high-error region and increase the sampling probability of data in these region for training the student model intuitively makes sense to me, which is essentially resembles to hard-negative mining in the classic machine learning literature.\n\n2. The experiments were well-executed and supports the main claim in the paper.\n\n3. The math on GFlowNets and their connection to amortized inference is helpful, especially helps contextualize the significance of the contributions.", "weaknesses": "1. The idea is not new; it closely resembles hard negative mining (i.e., sampling negative examples where the model shows high error), which limits the novelty of the proposed approach.\n\n2. While the idea of sampling more in high-error regions seems intuitively reasonable, its effectiveness may depend on whether the student model has sufficient capacity to fit the distribution. Also, I would like to see more comparisons and discussion with the active learning literature, such as uncertainty sampling, etc.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DihJZQUKP0", "reviewer": "Reviewer_y6wN", "summary": "The paper presents a novel method to improve amortized inference for complex distributions using an adaptive \"Teacher-Student\" training framework. The Student is an amortized sampler parameterized as a generative flow network (GFlowNets) and trained using RL. The primary contribution of the work is introducing the Teacher as an auxiliary  model that acts as the 'exploration policy' for the student. It is trained to guide the Student training by focusing on high-loss regions, thereby promoting the discovery of unexplored modes of the target distribution. The proposed method is evaluated on synthetic environments, diffusion-based sampling tasks, and biochemical discovery tasks, demonstrating improved mode coverage and sample efficiency.", "questions": "same as weaknesses", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "**Strengths** : \n- Addresses an important problem: Mode coverage/exploration is an important problem in the training of GFlowNets. The paper proposes a novel and interesting solution to the problem.\n\n- Very Well written :  I really enjoyed reading the paper. The paper did an excellent job of introducing and walking through the relevant literature and the methods and putting itself in context. Although I was not myself very familiar with the specific work line of work around GFlowNets, I was easily able to follow along all the details. \n\n- Lots of interesting details : The training formulation was interesting, especial given I wasn't very familiar with the glownets literature before this. I also particularly liked the use of a search procedure combined with the Teacher network to reduce the teacher network induced bias in the exploration process (although this was already introduced in previous work!). This approach effectively guides the student towards more diverse solutions, improving the overall learning efficiency and robustness. \n\n- Impressive empirical results: The paper demonstrates the versatility of the Teacher-Student framework by applying it to a range of tasks, including synthetic benchmarks and biochemical discovery problems. The empirical results consistently show that the Teacher-Student setup leads to better mode coverage and training efficiency compared to existing methods. Especially the results in more complicated tasks with a large number of modes.", "weaknesses": "**Weaknesses/Questions**\n-  The introduction of an adaptive Teacher adds additional complexity to the training process, requiring the joint optimization of both Teacher and Student networks. At least in the RL literature, these types of exploration methods were tried and given up on as they required extensive tuning and didn't scale well enough. I'm curious how the authors think that compares with the use cases here and if the authors genuinely believe the results shown in the paper will hold the test of time?\n\n- This is maybe a dumb question. But I do wonder how these methods compare with using standard sampling based strategies for mode discovery e.g  thompson sampling etc. My understanding is those become intractable as the problem size increases. The approach suggested seems pretty complex and I do wonder if the those standard mode discovery methods could've helped make things simpler. \n\n- Details on the biochemical discovery experiments were a little unclear. Eg. How do you define the reward function etc was not very clear to me.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "my4Wqjt3tJ", "reviewer": "Reviewer_NTRq", "summary": "This work focuses on the efficient exploration of RL training during amortized inference. The primary contribution lies in developing an adaptive training distribution to guide the amortized sampler in prioritizing difficult ones. The proposed method is examined in a collection of benchmarks, including both synthetic and real-world scenarios. The exploration efficiency and other benefits are reflected in both mode coverage and sample efficiency.", "questions": "See Weakness part", "limitations": "", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "I can easily follow this work, and this work tries to amortize prediction by simply conditionally inputting some variables.\n\nOverall, (1) I find this work easy to follow with clear motivations. Decision-making for amortized inference, particularly the development of GFlowNets, is impactful, and this work focuses on an important issue, namely efficient exploration under RL frameworks in the field. (2) The developed strategy is novel and practical in implementation. (3) The experiments are inspiring and well-supported claims.", "weaknesses": "While a lot of merits in this work, I find some parts are necessary to modify or revise.\n\n---\n\n(1) It seems to lack the necessity of amortized inference. In line28-30, it states the mechanism of amortized inference and related bottleneck. It is necessary to include the role of amortized inference compared with traditional methods such as MCMC, e.g., citing [1] and adding something like\n\n\"The amortized inference adopts a shared inference module for all data points instead of performing inference one by one. In this way, we can reuse the computational module for other data point's inference.\"\n\n(2) There exists literature work [2] that raises similar learning modules in terms of training adaptative distributions for few-shot experimental design; it is necessary to discuss them in detail in Section 4 related work.\n\n\n(3) Other suggestions or questions: (i) Figure 1 is clear enough, but I am not sure whether there should exist links between the Buffer and the Teacher or the Student to reveal the data flow. (ii) in Line-74, it says \"we believe that trajectories with high loss are particularly informative for mode coverage\", is this a hypothesis? Are there any explanations from either experiments or other intuitions? (iii) In Line 60, it says \"the Student's target distribution does not depend on the Teacher\", hence I am wondering whether the optimization pipeline is an adversarial game. (iv) In line 237, I want to know how to balance the student, the teacher or a buffer.\n\n\n**Reference:**\n\n[1] Margossian C C, Blei D M. Amortized Variational Inference: When and Why?[J]. arXiv preprint arXiv:2307.11018, 2023.\n\n[2] Wang, Cheems, et al. \"Robust Fast Adaptation from Adversarially Explicit Task Distribution Generation.\" arXiv preprint arXiv:2407.19523 (2024).\n\n---\n\nI will be happy to update my score if these concerns are well addressed during the rebuttal discussion.\n\n---\n\nPost Rebuttal\n\nThe author has well addressed my concerns, and I updated my score to accept.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "7WkcJH82L9", "reviewer": "Reviewer_CGLL", "summary": "This paper introduces a method for training a neural network to approximate a distribution with a specified unnormalized density. Specifically, it proposes an adaptive training distribution, termed the \"Teacher,\" which guides the primary amortized sampler, or \"Student,\" by prioritizing high-loss regions.", "questions": "1. Convergence Rate of Algorithm 1: What is the convergence rate of your Algorithm 1? Although it performs well in exploring more modes in the example tasks, the convergence rate is also a key factor in sampling methods. Could you provide more details on this?\n\n2. Selection of Behavior Policy: How did you determine the behavior policy during training (line 220)? From lines 827 to 831, it appears that different tasks require different ratios. What standard guided these choices? The explanation on lines 835 to 840 lacks proof of the algorithm’s robustness across different ratios, which raises concerns about whether these choices were made deliberately and may impact the generality of the results. The same issue applies to the choice of C in Table 5 and α in Tables 6 and 7.\n\n3. Choice of Backward Policy: How did you select the backward policy during training? It appears that a uniform random policy is used in the deceptive grid world and biochemical design tasks. In my understanding, the backward policy in your algorithm plays a role similar to the proposed transition kernel in MCMC, which is crucial for convergence. Could you elaborate on the role of the backward policy in your algorithm and its impact on the convergence rate of Algorithm 1 across different tasks?\n\n4. Algorithm Performance on High-Dimensional Tasks: How does the algorithm perform in high-dimensional tasks? Sampling from a specified unnormalized density is particularly challenging in high dimensions. Testing the algorithm on high-dimensional tasks could strengthen the evaluation of its effectiveness.\n\n5. Performance in the Manywell Task: In the Manywell task, the performances of PER+LS and the Teacher-based methods appear similar (table 8). What is the underlying intuition for this?\n\nI would be willing to increase the score if the above questions could be well clarified.", "limitations": "", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The Teacher model has the potential to generalize across the Student's high-loss regions, thereby improving mode coverage. The algorithm’s effectiveness is demonstrated through both discrete and continuous tasks, comparing favorably against other GFlowNet algorithms.", "weaknesses": "This paper appears to fall slightly below the standard expected at ICLR for several reasons:\n\n1. A critical issue is the lack of a solid theoretical foundation; the paper primarily reports numerical results without a deeper mathematical analysis. For instance, there is no mathematical description or guarantee of convergence rate for Algorithm 1.\n\n2. Regarding the experiments, the explanation of the architecture design and the choice of hyperparameters would benefit from greater clarity and justification. I will outline these concerns in more detail in the questions below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)