AIRI and Skoltech: ENOT Generative Model Training Method

Main article: Generative artificial intelligence

2024: Introduction of a method to accelerate training of generative models up to 10 times

Scientists from the AIRI Institute and Skoltech have proposed an approach to solving an extensive class of complex computational problems based on Optimal Transport (OT), applicable in machine learning and mathematical modeling. The method will speed up training of models from 3 to 10 times. AIRI announced this on December 5, 2024.

Optimal transport techniques are increasingly being used in training generative models to synthesize artificial data, such as images or texts. Another significant application is the adaptation of models to data from new sources, which is especially relevant in medicine, where work is often associated with small and disparate samples. However, existing methods of solving OT problems using neural networks face a number of problems, such as high instability of training and the need for complex intermediate transformations.

The key advantage of the method proposed by scientists, implemented on the JAX framework and called ENOT, was the introduction of entropy regularization. This made it possible to achieve significant acceleration of calculations - from 3 to 10 times - and improvement of target metrics of model performance. Initially, experiments were carried out on two-dimensional data, and later the method was tested on image generation tasks, style transfer and reconstruction of three-dimensional objects, which confirmed its versatility.

As for practical application, the tasks of optimal transport are multidisciplinary, so it can be applied in a variety of fields. In particular, we applied it to imitation training - when an expert shows certain actions, the agent tries to imitate the behavior, and the system evaluates how similar the agent's actions are to those of the expert. An example is a dance lesson, when a teacher shows movement, and a student tries to repeat it, "explained Nazar Buzun, head of the Teaching Representations group of the Strong AI in Medicine laboratory at the AIRI Institute.

Despite the theoretical format of the article, the method is based on an intuitive idea - it is proposed to "pull" the generated to the expected one. It seems to me that today there are too many "black boxes" in our area. Of course, tenfold acceleration is a weighty argument, but I think that NeurIPS reviewers liked our method precisely for its intuition, "said Dmitry Dylov, director of the Strong AI in Medicine Laboratory at the AIRI Institute and associate professor at the Skoltech AI Technology Center.