Learning Energy Decompositions for Partial Inference of GFlowNets (2310.03301v1)
Abstract: This paper studies generative flow networks (GFlowNets) to sample objects from the Boltzmann energy distribution via a sequence of actions. In particular, we focus on improving GFlowNet with partial inference: training flow functions with the evaluation of the intermediate states or transitions. To this end, the recently developed forward-looking GFlowNet reparameterizes the flow functions based on evaluating the energy of intermediate states. However, such an evaluation of intermediate energies may (i) be too expensive or impossible to evaluate and (ii) even provide misleading training signals under large energy fluctuations along the sequence of actions. To resolve this issue, we propose learning energy decompositions for GFlowNets (LED-GFN). Our main idea is to (i) decompose the energy of an object into learnable potential functions defined on state transitions and (ii) reparameterize the flow functions using the potential functions. In particular, to produce informative local credits, we propose to regularize the potential to change smoothly over the sequence of actions. It is also noteworthy that training GFlowNet with our learned potential can preserve the optimal policy. We empirically verify the superiority of LED-GFN in five problems including the generation of unstructured and maximum independent sets, molecular graphs, and RNA sequences.
- Rudder: Return decomposition for delayed rewards. Advances in Neural Information Processing Systems, 32, 2019.
- Why is tanimoto index an appropriate choice for fingerprint-based similarity calculations? Journal of cheminformatics, 7(1):1–13, 2015.
- Survey of variation in human transcription factors reveals prevalent dna binding changes. Science, 351(6280):1450–1454, 2016.
- Flow network based generative models for non-iterative diverse candidate generation. Advances in Neural Information Processing Systems, 34:27381–27394, 2021a.
- Gflownet foundations. arXiv preprint arXiv:2111.09266, 2021b.
- Ludwig Boltzmann. Studien uber das gleichgewicht der lebenden kraft. Wissenschafiliche Abhandlungen, 1:49–96, 1868.
- Exploration by random network distillation. In International Conference on Learning Representations, 2019. URL https://openreview.net/forum?id=H1lJJnR5Ym.
- Reinforcement learning with trajectory feedback. In Proceedings of the AAAI conference on artificial intelligence, volume 35, pp. 7288–7295, 2021.
- Learning guidance rewards with trajectory-space smoothing. Advances in Neural Information Processing Systems, 33:822–832, 2020.
- Neural message passing for quantum chemistry. In International conference on machine learning, pp. 1263–1272. PMLR, 2017.
- Biological sequence design with gflownets. In International Conference on Machine Learning, pp. 9786–9801. PMLR, 2022.
- When to trust your model: Model-based policy optimization. Advances in neural information processing systems, 32, 2019.
- Junction tree variational autoencoder for molecular graph generation. In International conference on machine learning, pp. 2323–2332. PMLR, 2018.
- Algorithmic framework for model-based deep reinforcement learning with theoretical guarantees. In International Conference on Learning Representations, 2018.
- Learning gflownets from partial episodes for improved convergence and stability. In International Conference on Machine Learning, pp. 23467–23483. PMLR, 2023.
- Trajectory balance: Improved credit assignment in gflownets. Advances in Neural Information Processing Systems, 35:5955–5967, 2022a.
- Gflownets and variational inference. arXiv preprint arXiv:2210.00580, 2022b.
- Better training of gflownets with local credit and incomplete trajectories. arXiv preprint arXiv:2302.01687, 2023a.
- Generative augmented flow networks. In International Conference on Learning Representations, 2023b. URL https://openreview.net/forum?id=urF_CBK5XC0.
- Learning long-term reward redistribution via randomized return decomposition. In International Conference on Learning Representations, 2022. URL https://openreview.net/forum?id=lpkGn3k2YdD.
- Proximal policy optimization algorithms. arXiv preprint arXiv:1707.06347, 2017.
- Towards understanding and improving gflownet training. In Proceedings of the 40th International Conference on Machine Learning, Proceedings of Machine Learning Research. PMLR, 2023.
- Mastering the game of go with deep neural networks and tree search. nature, 529(7587):484–489, 2016.
- Dual policy iteration. Advances in Neural Information Processing Systems, 31, 2018.
- Reinforcement learning: An introduction. MIT press, 2018.
- Design-bench: Benchmarks for data-driven offline model-based optimization. In International Conference on Machine Learning, pp. 21658–21676. PMLR, 2022.
- Let the flows tell: Solving graph combinatorial optimization problems with gflownets. arXiv preprint arXiv:2305.17010, 2023.