Artificial Intelligence and Machine Learning

Teaser

AI systems that recognise images or translate text are built from neural networks. How such a network is put together has a big influence on how well it works. Finding a good structure usually means building and training many versions, which takes a lot of computing time. Researchers in the group of Prof. Volker Tresp at LMU Munich and the Munich Center for Machine Learning (MCML) developed DiffNAS, a method that uses generative AI to suggest network structures and quick tests to sort out weak ones early. On a standard benchmark, DiffNAS found the best available structures with very little search time. The experiments used more than 100,000 node hours on the Hawk supercomputer at HLRS and additional time on its successor Hunter.

The Challenge

AI systems that recognise faces, translate text or detect diseases in medical images are based on neural networks. A neural network processes information in many small steps. Which steps it uses, and how they are connected, is called its architecture. The architecture has a large influence on how well the network works.

Finding a good architecture is difficult. There are thousands or even millions of possible combinations. The usual way to find out whether one of them is good is to train it with data and then test it. For a single network this can take hours or days, so testing many candidates quickly becomes very expensive.

Methods that search for good architectures automatically can help. Many of these methods use a second AI model that has learned to predict how good a network will be. To build this predictor, however, a large number of networks must first be trained, so the cost does not go away.

The Approach

A team in the group of Prof. Volker Tresp at LMU Munich and MCML developed DiffNAS (Diffusion Neural Architecture Search) to search for good architectures without such a predictor.

DiffNAS draws each architecture as a small diagram of boxes and arrows and uses a diffusion model to create new diagrams. Diffusion models are the technology behind many AI image generators. They learn by taking examples, mixing them up step by step until only noise is left, and practising how to reverse this. After training, they can start from random noise and build something new. In DiffNAS, the new result is the design of a neural network.

To steer the model towards good designs, DiffNAS uses quick tests that work on networks that have not been trained yet. A single image is sent through the untrained network, and the test measures how varied the signals inside the network are. This takes only a moment and gives a useful hint about how well the network might perform after training. DiffNAS uses these test scores while the diffusion model learns, rewards the model for designs with high scores, and finally uses the scores to pick the most promising designs (Figure 1).

Why Supercomputing Was Needed

Once DiffNAS is trained, a search takes very little time. Developing and testing the method, however, required a great deal of computing power. The diffusion models had to be trained and fine-tuned many times. The method was tested on several image collections and on two different types of networks. In further tests, individual parts of the method were switched off one at a time to see how much each part contributes.

The project used more than 100,000 node hours on the Hawk supercomputer at the High-Performance Computing Center Stuttgart (HLRS), as well as time on its successor Hunter. Many of these calculations are independent of each other, so they could run at the same time on different parts of the machine. Published benchmark data, which already contain the test results of thousands of trained networks, meant that these networks did not have to be trained again.

Key Findings

The main test used NAS-Bench-201, a public collection of 15,625 network architectures whose test results are already known. In two image-recognition tasks, the networks sort small photos into 10 categories such as "cat" or "ship" (CIFAR-10), or into 100 categories (CIFAR-100). In both tasks, DiffNAS found the best architecture in the collection, with 94.37 and 73.51 percent of the images classified correctly. DiffusionNAG, a leading method that also uses a diffusion model, reaches the same result, but it needs a trained predictor. The search itself took DiffNAS about 15 seconds on one graphics processor, about 17 times less than DiffusionNAG. This does not include the time needed to train the diffusion model beforehand.

On two further image collections, with photos of aircraft and of pets, the results were similar to or slightly below those of DiffusionNAG (Table 1). The generated designs are also good on average: when DiffNAS created 1,000 designs for CIFAR-10, they reached 93.4 percent on average, compared with 87.1 percent for 1,000 randomly chosen designs. When the quick tests or the reward step were switched off, the generated designs became clearly worse.

On MobileNetV3, a family of compact networks built for smartphones, DiffNAS did not yet reach the accuracy of DiffusionNAG. Adapting the quick tests to such networks is the next step.

Who Benefits

The results are useful above all for researchers who build tools that design AI models automatically. DiffNAS shows that generative AI and quick tests on untrained networks can work together. Methods like DiffNAS could help universities and companies with small computing budgets to find good network designs with far fewer training runs.