HSE Scientists Optimise Training of Generative Flow Networks

© iStock

Researchers at the HSE Faculty of Computer Science have optimised the training method for generative flow neural networks to handle unstructured tasks, which could make the search for new drugs more efficient. The results of their work were presented at ICLR 2025, one of the world’s leading conferences on machine learning. The paper is available at Arxiv.org.

Generative Flow Networks (GFlowNets) are a class of machine learning algorithms that build complex objects step by step. Researchers use them to search for new proteins and drugs, and to optimise transport systems. 

For GFlowNets to discover such complex structures, researchers specify the desired properties of the target object. The closer the network’s proposed solution is to these properties, the higher the reward it receives. GFlowNets aim to solve problems in a way that maximises their reward. They do not rely on data directly, but instead on the reward, which is computed using an equation known as the value function.

The search for a complex object can be compared to assembling a Lego model, where pieces are added step by step until the object is complete, with each model assigned a specific value—for example, a plant model might be valued higher than an animal model. Unlike other machine learning methods that would strive to construct a plant at any cost, GFlowNets generate a variety of objects—but plants more frequently than animals—because the reward for plants is higher.

In this type of search, GFlowNets rely on two stochastic policies that operate together: a forward policy and a backward policy. The forward policy can be thought of as a construction foreman, deciding the next step and estimating the probability of the subsequent state, while the backward policy acts as a deconstruction expert, identifying the preceding step. Maintaining balance between these flows is crucial but difficult to achieve. First, it requires significant computing power. Second, backward policies lack flexibility: researchers usually prevent them from adapting during the search or from observing the actions of the forward policy.

HSE scientists have developed a way to optimise backward policies using a method called Trajectory Likelihood Maximisation (TLM). They refined the backward policy’s algorithms so that it can be continuously checked against the steps of the forward policy.

Timofey Gritsaev

'We designed the search for the optimal solution to resemble a negotiation, where both sides are ready to adjust their positions. In highly uncertain problems, the backward policy serves only as an auxiliary tool that improves the results of the forward policy. Our goal was to make the backward policy more flexible, and we finally succeeded,' explains Timofey Gritsaev, co-author of the paper and Research Assistant of the Centre for Deep Learning and Bayesian Methods at the HSE FCS AI and Digital Science Institute.

After implementing TLM, the reward function that measures the backward model’s success became more complex. Nevertheless, despite this increased complexity, the overall search system became faster and more efficient.

Nikita Morozov

'Our method explores the space of possible solutions noticeably faster and identifies more high-quality options. Overall, this approach brings generative models closer to reinforcement learning methods,' explains Nikita Morozov, Junior Research Fellow of the Centre for Deep Learning and Bayesian Methods at the AI and Digital Science Institute of the HSE FCS.

The authors of the study are confident that their work will benefit specialists using GFlowNets across various fields, including the search for new medicinal compounds, the development of materials with specific properties, and the fine-tuning of large language models. Thanks to these networks’ ability to efficiently explore vast solution spaces and quickly identify the best options, the demand on computing power can be significantly reduced.