Adaptive teachers for amortized samplers

Thank you for your detailed feedback and effort in reviewing our paper. We have attempted to answer your questions, provide new experimental results to support our claims, and clarify some possible misunderstandings below.

W1 (lack of theoretical analysis)

Please see Q1 below.

W2 (lack of explanation of the architecture design and hyperparameters)

As you mentioned, the explanation of the architecture design is crucial for greater clarity and justifiction. We mention details on implementation for each experiment in Appendix C. Here we explain more details on model architecture for each experiment.

Deceptive Grid World: We use a two-layer MLP with 256 hidden units for the parameterized policy $P_F(\cdot;\theta)$ along with a learnable parameter for $\log Z_\theta$. The backward policy $P_B$ is fixed to a uniform policy.

Diffusion Sampling: We employ the same architecture as [1, 2]. We encode the diffusion timestep $t$ with 128-dimensional harmonic (Fourier) features use 2 linear layers with $N=64$ hidden units to extract the signal ($x$) feature. We use a two-layer MLP with $N$ hidden units to extract feature for $x_t$ and concatenate it with the signal feature. Finally, we apply a three-layer MLP with $N$ hidden units to this concatenated representation to get $u(x_t, t;\theta)$. We initialize $\log Z_{\theta}$ to 0 for all methods. For the Manywell task, we increase $N$ from 64 to 256 to accommodate the high-dimensional tasks, and apply this adjustment to all baselines.

Biological and Chemical Discovery: We use a similar setting to that proposed by [3]. To parameterize the forward policy, we adopt a relative edge flow policy parametrization mapping (SSR) from [3]. For QM9 and sEH tasks, we employ a two-layer architecture with 1024 hidden units, while for the other tasks, we choose to use a two-layer architecture with 128 hidden units. We initialize $\log Z_{\theta}$ to 5.0 for all methods. For the backward policy, we use a fixed uniform policy.

[1] Sendera, Marcin, et al. "Improved off-policy training of diffusion samplers." The Thirty-eighth Annual Conference on Neural Information Processing Systems.

[2] Zhang, Qinsheng, and Yongxin Chen. "Path Integral Sampler: A Stochastic Control Approach For Sampling." International Conference on Learning Representations.

[3] Shen, Max W., et al. "Towards understanding and improving GFlowNet training." International Conference on Machine Learning. PMLR, 2023.

Concerns about hyperparameters are addressed in the questions below.

Q1 (convergence rate)

Our paper makes an methodological and empirical contribution. We believe that our numerical results are sufficient to validate our proposed idea, which we expect to be influential for machine learning researchers interested in the practical aspect of GFlowNets, e.g., their application to drug discovery [1], biological sequence design [2], and language models [3,4]. To show its wide applicability, we tested the proposed algorithm on 10 tasks from different domains and found improvements over exploration methods from past work in a wide range of setting.

Nevertheless, we appreciate your emphasis on the importance of convergence rates in evaluating sampling algorithms. Providing a theoretical convergence rate for Algorithm 1 is indeed valuable; however, it presents significant challenges due to the complexity inherent in deep learning-based methods. Convergence rates of deep learning models even under a fixed sampling distribution remains an difficult problem in the field, not to mention RL agents with a dynamic behavior policy updated in a bi-level optimization procedure.

[1] Shen, Tony, et al. "TacoGFN: Target-conditioned GFlowNet for Structure-based Drug Design". Transactions on Machine Learning Research, 2024. [2] Jain, Moksh, et al. "Biological sequence design with gflownets." International Conference on Machine Learning (ICML), 2022. [3] Hu, Edward J., et al. "Amortizing intractable inference in large language models." International Conference on Learning Representations (ICLR), 2024. [4] Lee, Seanie, et al. "Learning diverse attacks on large language models for robust red-teaming and safety tuning." arXiv preprint arXiv:2405.18540, 2024.

Dear Reviewer CGLL,

As we reach the end of the discussion period, we'd like to ask if we can provide any more information that could affect your assessment of the paper. In the discussion period, we believe that we clarified the concerns and questions raised. Specifically, we included additional experiments to address concerns about scalability and flexibility under different backward policy settings (the learned $P_B$). This will be a good improvement to our manuscript; thanks again for your feedback.

The authors

Thank you for your review and the positive assessment of our paper. We answer your concerns below.

Why amortized inference?

Thank you for your suggestion. We agree that the ability to reuse a shared computational module for inference across multiple data points, as opposed to performing inference independently for each data point, is a major motivation for amortized inference compared to MCMC. Following your recommendation, we have cited the suggested reference and revised the first paragraph of introduction to better highlight this key advantage of amortized inference.

Related work on experimental design

Thank you for pointing out this relevant literature. We agree that it is related to our work, as our Teacher plays the same role as the entropy regularized adversary that generates tasks. We have included discussion of the suggested reference in the related work (Sec. 4).

Figure 1

The illustration is intended to show that the behavior policy (Teacher, Student, or Buffer) contributes to the data flow for both Teacher and Student training. This implies that the Buffer sometimes provides data to the Teacher and Student during training.

In the revised version, we have added an arrow from the Student to the Buffer to Figure 1, as trajectories are collected from both models.

Why are high-loss trajectories informative?

This is a hypothesis, motivated by the following arguments:

• In all learning systems, samples with high error, where the current iteration of the model struggles, tend to be more informative. This is a principle used in active learning, hard example selection, curriculum learning, prioritized experienced replay in RL, etc.
• For amortized inference systems in particular, discovery of poorly modeled modes (especially those whose density is underestimated) is critical. This is because errors in already observed modes can self-correct, as they are revisited during training, but missing modes are unlikely to be visited by on-policy (or near-on-policy) sampling, making them harder to recover.

Thus, to promote discovery of modes, the Teacher should guide the Student toward samples where the sampling density divergences from the target, especially favouring the samples whose density is underestimated. This motivates the proposed reward for the Teacher (equation 5).

Is this an adversarial game?

Good question. Assuming the function classes of the Student and Teacher policies can express the unique point where both achieve zero loss, the joint learning problem between Teacher and Student is not adversarial in the sense of having a saddle point at the optimum. The losses for both models are strictly positive and are zero precisely at the optimum (cf. Proposition 1 in Appendix A), so any deviation from the optimum will not decrease the losses of teacher, student, or both.

Furthermore, the Student's reward does not depend on the parameters of the Teacher, even though its training policy is, which implies that the optimality of a Student policy is independent of the Teacher's parameters.

However, if the optimal policies are not representable by the policy networks of the Student and Teacher, a saddle point at the optimum is possible.

Details of behavior policies

The details are provided in Appendix B (lines 827-831). The balancing rule is straightforward: the Teacher focuses on multi-modal exploration, the Student emphasizes exploitation, and the replay buffer is used for sample efficiency. Depending on the characteristics of the target task, users can adjust the balance according to their specific needs, similar to how exploration-exploitation trade-offs are tuned in other exploration methods in RL.

We appreciate your valuable input; please let us know if we can provide any further clarifications.

Thanks for the author's detailed response and clarifications. After revision, the updated manuscript is complete enough. Taking other reviewers' comments and my assessments, I think this work is well-motivated and proposes a novel approach to Bayesian deep learning with sufficient evaluation. Hence, I have updated my score and think this work deserves acceptance.

Thank you for such a positive assessment of our paper's relevance, exposition, and empirical evaluation. We've attempted to answer your questions below.

Complexity of training and simpler exploration strategies

The question of whether the benefits brought by the proposed algorithm are worth the added complexity of training a teacher model, relative to simpler exploration methods, is very important, and we appreciate the opportunity to comment upon this point.

Certainly, the additional tunable parameters are unavoidable when we introduce a new component (like the Teacher network) to an algorithm. However, these parameters are often beneficial in that they provide controllability of a desired behavior, in our case, exploration. This can be compared to the way that GANs improve the training of a generator by introducing a secondary network, the discriminator, into the optimization.

Although the joint optimization of two networks increases training complexity, there are reasons to believe the benefits outweigh the shortcomings, even at large scales. It is clear that large-scale tasks where exploration is difficult require a well-designed training policy for an agent. The more complex the task, the harder it is for a replay buffer to capture all modes of the target distribution. The Teacher network can be viewed as an amortized replay buffer that provides generalization abilities that a buffer cannot -- it can generate arbitrarily many new training samples for the Student on the fly. This approach is fundamentally more scalable compared to non-learned sample selection methods.

As for comparison to simpler exploration methods:

• We compared our method to simple approaches such as epsilon-greedy and replay buffer methods, but they did not achieve satisfactory performance, particularly in large-scale tasks.
• Methods like Thompson sampling, which rely on maintaining a Bayesian posterior over parameters, become extremely complex at scale, while RND requires training a second model, typically of a similar complexity to the sampler (just as our Teacher). Note that we also compare with both of these methods in the form they were proposed in prior work, but again found them to underperform.

Finally, we want to emphasize that methodological simplicity is not equivalent to the complexity or scalability of an algorithm. While you mention that our method requires joint optimization and may seem complex, this does not necessarily mean it introduces significant complexity at scale. Our approach simply involves training one additional network at any scale, leading to a constant multiplicative increase in number of parameters.

Thank you for your review. We appreciate that you described the proposed algorithm as intuitive and found the exposition and experiments clear and well-executed. Below we answer the questions and concerns you raised.

Novelty and resemblance to hard negative mining

Our main idea is to amortize the sampling of high-loss examples through the use of a teacher network, rather than selecting hard examples from a dataset. This clearly distinguishes our work from hard example mining [Robinson et al., 2021].

In the RL setting (of which GFlowNets are one instance), hard example mining is related to prioritized experience replay (PER), in which a buffer functions as the dataset from which hard examples are chosen. We indeed consider PER as a baseline in our study.

Model capacity

We agree that the effectiveness of the sampler resulting from our proposed algorithm depends on the student model's capacity to fit the distribution. This is a consideration common to all training methods for amortized samplers and involves such questions as the neural network architecture and the update rule (optimizer) used for the policy network for each given sample.

However, the problem we address is orthogonal: for a student model of given architecture, how do we best select the training distribution (i.e., behavior policy) that provides it the samples to learn from? Our solution, which introduces an auxiliary teacher model, leads to improved exploration relative to other techniques that use the same student model architecture. In fact, our technique decouples the modeling capacity of the student from the behaviour policy, unlike other methods (noisy on-policy, Thompson sampling, etc.) that use modifications to the student's policy to induce exploration.

Comparison with active learning

Thank you for pointing out the connection with active learning. We agree that there are strong connections between active learning and our method, as both aim to leverage information about regions where the current model performs poorly. However, there are also clear distinctions between the problems considered in the two areas.

• Active learning is primarily a framework for supervised learning, where the goal is to select the most informative data points for labeling. These methods typically rely on a form of predictive uncertainty to select inputs $x$ that optimally inform the learning of a mapping to labels $y$, using information derived from the classifier's probabilistic model $p(y\mid x)$ but without seeing the true label $y$. Some forms of uncertainty used to guide example selection include margin sampling, maximum-entropy sampling, etc., as well as Bayesian uncertainty quantification approaches (ensembles and BNNs, inter alia), all of which query the samples $x$ where the predictor is, in some sense, most uncertain of the label.
• In contrast, our method is based on reinforcement learning (RL), not supervised learning. Instead of training a classifier to model $p(y\mid x)$, we train a sequential decision-making agent characterized by a policy $P_F(\tau) = \prod_{t=1}^n p(s_t|s_{t-1})$, where the trajectory $\tau$ represents a sequence of states $(s_0, s_1, \ldots, s_n)$. Gradient updates are made using a loss that depends on a trajectory $\tau$, which is not necessarily sampled from the policy itself. Our solution to the trajectory selection problem does notes not rely on predictive uncertainty. Instead, it leverages information from the loss values, which are computed using the true terminal reward values.

The two problems have fundamentally different characteristics, despite their conceptual similarities: in active learning, one seeks areas with high uncertainty (that is, possibly high unknown loss value), while we amortize the sampling of areas with high (known) loss value.

In the revised manuscript (Sec. 4, Related Work), we have included a more detailed comparison with the active learning literature, including uncertainty sampling methods.

We believe there is a slight misunderstanding. We agree that diffusion models are very useful and that we must continue researching how to improve such models.

Diffusion models and GFNs are orthogonal: A GFN is a training method for diffusion models without data, but with energy/reward.

Training diffusion models with denoising score matching (DSM) or maximum likelihood estimation (MLE)—typical methods for diffusion models—is scalable and practical when massive datasets are available. However, when we aim to enhance diffusion models with a reward model or perform intractable inference using energy functions, we need reinforcement learning (RL)-like training methods because there is no such massive dataset to imitate (making DSM or MLE on a dataset impossible). Among these methods, GFN are a promising candidate.

For example, Venkatraman et al. [1] demonstrated that fine-tuning diffusion models with a reward model can be effectively applied to various tasks, including inverse image problems, language model infilling using discrete diffusion, text-to-image model fine-tuning, and offline RL through GFN training over diffusion models.

Moreover, Seong et al. [2] shows that using the same GFN objective with Venkatraman et al. [1], we can model molecular dynamics (MD) by sampling rare transition paths. In such scientific discovery applications, there are desired models called Boltzmann Generators that aim to sample proportionally to the Boltzmann energy distribution $e^{-E(x)/T}$ (in MD and N-body particle simulations). For Boltzmann Generator training, we believe this combination is particularly useful: (1) we have to use diffusion models, and (2) train them using a GFN-like objective. There is active research following (1) and (2) to meet these demands [3, 4].

[1] Venkatraman et al. "Amortizing Intractable Inference in Diffusion Models for Vision, Language, and Controls", NeurIPS 2024.

[2] Seong et al. "Collective Variable Free Transition Path Sampling with Generative Flow Network", ICML Workshop on Structured Probabilistic Inference and Generative Modeling 2024.

[3] Sendera et al. "Improved Off-policy Training of Diffusion Samplers", NeurIPS 2024.

[4] Akhound-Sadegh et al., "Iterated Denoising Energy Matching for Sampling from Boltzmann Densities", ICML 2024.