Learning to Scale Logits for Temperature-Conditional GFlowNets

Minsu Kim, Joohwan Ko, Taeyoung Yun, Dinghuai Zhang, Ling Pan, Woo Chang Kim, Jinkyoo Park, Emmanuel Bengio, Yoshua Bengio
Proceedings of the 41st International Conference on Machine Learning, PMLR 235:24248-24270, 2024.

Abstract

GFlowNets are probabilistic models that sequentially generate compositional structures through a stochastic policy. Among GFlowNets, temperature-conditional GFlowNets can introduce temperature-based controllability for exploration and exploitation. We propose Logit-scaling GFlowNets (Logit-GFN), a novel architectural design that greatly accelerates the training of temperature-conditional GFlowNets. It is based on the idea that previously proposed approaches introduced numerical challenges in the deep network training, since different temperatures may give rise to very different gradient profiles as well as magnitudes of the policy’s logits. We find that the challenge is greatly reduced if a learned function of the temperature is used to scale the policy’s logits directly. Also, using Logit-GFN, GFlowNets can be improved by having better generalization capabilities in offline learning and mode discovery capabilities in online learning, which is empirically verified in various biological and chemical tasks. Our code is available at https://github.com/dbsxodud-11/logit-gfn

Cite this Paper

BibTeX
@InProceedings{pmlr-v235-kim24s, title = {Learning to Scale Logits for Temperature-Conditional {GF}low{N}ets}, author = {Kim, Minsu and Ko, Joohwan and Yun, Taeyoung and Zhang, Dinghuai and Pan, Ling and Kim, Woo Chang and Park, Jinkyoo and Bengio, Emmanuel and Bengio, Yoshua}, booktitle = {Proceedings of the 41st International Conference on Machine Learning}, pages = {24248--24270}, year = {2024}, editor = {Salakhutdinov, Ruslan and Kolter, Zico and Heller, Katherine and Weller, Adrian and Oliver, Nuria and Scarlett, Jonathan and Berkenkamp, Felix}, volume = {235}, series = {Proceedings of Machine Learning Research}, month = {21--27 Jul}, publisher = {PMLR}, pdf = {https://raw.githubusercontent.com/mlresearch/v235/main/assets/kim24s/kim24s.pdf}, url = {https://proceedings.mlr.press/v235/kim24s.html}, abstract = {GFlowNets are probabilistic models that sequentially generate compositional structures through a stochastic policy. Among GFlowNets, temperature-conditional GFlowNets can introduce temperature-based controllability for exploration and exploitation. We propose Logit-scaling GFlowNets (Logit-GFN), a novel architectural design that greatly accelerates the training of temperature-conditional GFlowNets. It is based on the idea that previously proposed approaches introduced numerical challenges in the deep network training, since different temperatures may give rise to very different gradient profiles as well as magnitudes of the policy’s logits. We find that the challenge is greatly reduced if a learned function of the temperature is used to scale the policy’s logits directly. Also, using Logit-GFN, GFlowNets can be improved by having better generalization capabilities in offline learning and mode discovery capabilities in online learning, which is empirically verified in various biological and chemical tasks. Our code is available at https://github.com/dbsxodud-11/logit-gfn} }
Endnote
%0 Conference Paper %T Learning to Scale Logits for Temperature-Conditional GFlowNets %A Minsu Kim %A Joohwan Ko %A Taeyoung Yun %A Dinghuai Zhang %A Ling Pan %A Woo Chang Kim %A Jinkyoo Park %A Emmanuel Bengio %A Yoshua Bengio %B Proceedings of the 41st International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2024 %E Ruslan Salakhutdinov %E Zico Kolter %E Katherine Heller %E Adrian Weller %E Nuria Oliver %E Jonathan Scarlett %E Felix Berkenkamp %F pmlr-v235-kim24s %I PMLR %P 24248--24270 %U https://proceedings.mlr.press/v235/kim24s.html %V 235 %X GFlowNets are probabilistic models that sequentially generate compositional structures through a stochastic policy. Among GFlowNets, temperature-conditional GFlowNets can introduce temperature-based controllability for exploration and exploitation. We propose Logit-scaling GFlowNets (Logit-GFN), a novel architectural design that greatly accelerates the training of temperature-conditional GFlowNets. It is based on the idea that previously proposed approaches introduced numerical challenges in the deep network training, since different temperatures may give rise to very different gradient profiles as well as magnitudes of the policy’s logits. We find that the challenge is greatly reduced if a learned function of the temperature is used to scale the policy’s logits directly. Also, using Logit-GFN, GFlowNets can be improved by having better generalization capabilities in offline learning and mode discovery capabilities in online learning, which is empirically verified in various biological and chemical tasks. Our code is available at https://github.com/dbsxodud-11/logit-gfn
APA
Kim, M., Ko, J., Yun, T., Zhang, D., Pan, L., Kim, W.C., Park, J., Bengio, E. & Bengio, Y.. (2024). Learning to Scale Logits for Temperature-Conditional GFlowNets. Proceedings of the 41st International Conference on Machine Learning, in Proceedings of Machine Learning Research 235:24248-24270 Available from https://proceedings.mlr.press/v235/kim24s.html.

Related Material