HSE, AIRI and Sber AI: Diffusion Neural Network Model

Main article: Neural networks (neural networks)

2024: Announcement of the diffusion neural network model

Scientists from the Center for Artificial Intelligence and the Faculty of Computer Science of the Higher School of Economics, as well as the Institute of Artificial Intelligence AIRI and Sber AI have developed a diffusion model structure for which it is possible to set eight types of noise distribution. Instead of the classical structure of the Markov chain model and the application of normal distribution, scientists proposed a star-shaped model where it is possible to choose the type of distribution. This will help solve problems in different geometric spaces using diffusion models. This was announced on February 15, 2024 by representatives of the Higher School of Economics.

As reported, over the past 20 years, generative neural networks have begun to work better. If earlier they created not very high-quality texts and images in one step, then with the advent of diffusion models - a type of generative neural networks - the process became gradual, and the result improved.

Diffusion neural networks are based on the probability model of noise cancellation and diffusion, or DDPM. The model works like this: at each stage, random changes are added to the data. For example, colors or brightness may change with each step. These changes gradually reduce noise and make the data more similar to the desired result until the chaos produces a final image.

The model is based on the Markov chain, which gradually adds noise, and then also gradually reverses the diffusion process in order to obtain the original data, for example, a picture with a cat. The neural network learns these transformations from training data, which contains an example of the original image and its noisy versions.

𝓲

cvpr2022-tutorial-diffusion-models.github.io

Such models generate pictures, sounds, but they cope worse with more complex tasks, for example, generating volumetric structures. This is due to the fact that the noise steps of the diffusion model only work using the normal distribution. And if the source objects have constraints, they cannot be set and saved throughout the steps.

The team of researchers proposed a type of model that optimizes the process of working with data. In this diffusion model structure, it has become possible to change the type of noise distribution. To achieve this, the researchers transformed the structure of the model into a star-shaped one, where all the states were not inside the Markov chain, but diverged from the original object to the sides.

For example, the task of a neural network is to generate a molecule. There are three types of atoms in the molecule that are defined using discrete data. If you noise this data with a normal distribution, then the types of atoms will take on values ​ ​ that do not exist in the real world. In a star-shaped model, we can choose the desired type of distribution, in which the data will not be distorted. told Andrey Okhotin, trainee researcher at the Center for Deep Learning and Bayesian Methods of the Institute of Artificial Intelligence and Digital Sciences, Faculty of Computer Science, Higher School of Economics

There are two components in the model structure. The first is responsible for noise of the object by step-by-step removal of information, and the second learns to take a step back in this chain. You can define a model for eight types of distributions that support data constraints.

We have moved on to the structure of the reverse process. If earlier each next state could be obtained using only one previous state, now each state of the object depends on all previous ones. With this structure, the information is collected into one object, which we called tail statistics, and fed into the neural network so that it takes the next step. This allows you to train the model more efficiently. told Dmitry Vetrov, scientific director of the Institute of Artificial Intelligence and Digital Sciences, FKN HSE, scientific consultant AIRI

Scientists compared the efficiency of the star-shaped model with classical diffusion models. On the tasks of generating text in normal mode, the model of scientists worked at the same level of quality. And in accelerated mode (with fewer generation steps), the model for images worked better and generated a dataset closer to the original.

With complex problems related to the generation of points in different geometric spaces - the sphere, simplex and matrix space describing ellipses - the star-shaped model coped better than the classical diffusion model.

source = HSE

In the problem of generating points on the sphere of the model, it was necessary to learn how to mark points in those places where, according to the 2020 geodetic dataset, fires most often occurred on the surface of the Earth. After that, the points that were in reality and those that generated the model were compared. The model generated points as close as possible to the original. The results obtained are comparable to existing methods for solving this problem.

source = HSE

In this paper, we proposed a more universal diffusion model that allows the generation of objects of complex structure. This will help apply such methods to a wider class of problems from the natural sciences, for example, biology, physics, chemistry, where there are structural limitations in the generation of objects: molecules, particle states, chemical compounds. reported by Aibek Alanov, Associate Fellow at the Center for In-Depth Learning and Bayesian Methods, Institute of Artificial Intelligence and Digital Sciences, FKN HSE, Research Fellow at AIRI

The study was supported by a grant for research centers in the field of artificial intelligence provided by the Analytical Center under the Government of the Russian Federation.