Accelerating neural network optimization: The power of second-order methods

Many of today’s AI applications, ranging from self-driving cars to large language models, are powered by neural networks that process huge amounts of data and recognize patterns in them. Optimization is a key part of the process of training these systems, during which scientists use techniques to improve a network’s performance by adjusting its parameters. Effective optimization methods are extremely important because they determine the efficiency and success of training — and, in the end, the usefulness of a system.

There are, however, many possible approaches to optimization and choosing the best technique for a specific network depends on several factors, including the task it will be used for, the amount of data available for training and the resources — both in terms of time and money — that are available for optimization. Scientists are therefore always working to develop new methods that can make optimization better and faster.

A team of scientists from the Mohamed bin Zayed University of Artificial Intelligence and other institutions recently presented a new approach for optimizing neural networks at the 38th Annual Conference on Neural Information Processing Systems (NeurIPS). It uses what are known as second-order methods to solve optimization problems related to variational inequalities, which are common in machine learning.

In the study, the researchers demonstrate that for the class of monotone inequalities with inexact second-order derivatives they examine, no faster second-order, or therefore, first-order, methods can theoretically exist. They also support this theoretical finding with experiments on a dataset, illustrating the practical applicability and effectiveness of their approach. The team’s findings suggest that their innovation has the potential to play a role in reducing the cost of optimization, particularly for large and complex neural networks.

The need for optimization

Optimization problems are everywhere, not only in machine learning, explains Artem Agafonov, a doctoral student in machine learning at MBZUAI and co-author of the study. “In some sense, our lives are optimization problems since we are always trying to figure out how to make the right decision to further our goals, whatever they may be,” he says. “We are doing something similar with neural networks.”

A goal of optimization is to minimize what is known as a loss function, which measures the difference between outputs predicted by a network and the actual data. This is essentially a measure of how poorly a system performs. Scientists use algorithms that iterate on the parameters of the network to reduce the loss until it stops changing significantly, a process known as convergence.

A common optimization technique is gradient descent, which seeks to minimize the loss function by iteratively updating parameters of a network in the direction that leads to the steepest decrease in the loss. Gradient descent is considered a first-order method because it relies on the first-order derivative, or gradient, of the loss function to move in the direction that reduces the loss. While gradient descent is a straightforward and widely used approach that can be effective in some cases, it can be extremely slow and ineffective in others.

Introducing second-order methods

Second-order methods are typically faster than first-order methods like gradient descent because they use not only the gradient of a function but also other information, like its curvature, which describes its rate of change. This allows a method to make more informed, often larger, steps towards reducing a loss function.

Second-order methods, however, can be computationally expensive for optimizing neural networks since the number of parameters in these models is huge, explains Dmitry Kamzolov, a research associate at MBZUAI and co-author of the study.

Agafonov, Kamzolov and their colleagues’ approach, which they call variational inequalities under Jacobian inexactness (VIJI), is powerful because of its ability to use second-order information without the heavy computational cost that has traditionally come with doing so. They achieve this by introducing approximations of the second-order derivatives, which makes the process computationally feasible without sacrificing performance.

“We found that we can use a cheap approximation of curvature information and it will give you almost the same convergence speed as calculating it exactly,” Agafonov says.

In their study, the researchers made second-order methods more practical by managing a concept known as Jacobian inexactness, which refers to the challenge of estimating the second-order derivatives of a function. “We store these gradients and create some inexact approximation of the second-order derivative,” Kamzolov says. “We don’t forget about the previous gradients but calculate some additional information from them and enhance the performance of the network because of this.”

At the core of their approach is a technique known as quasi-Newton approximation, which helps approximate the second-order derivative without needing to compute it exactly. This involves using previous gradients and updating them to get an approximation of the curvature, allowing the algorithm to increase the speed of convergence while keeping computational costs low.

Real-world impact and future research

Kamzolov explains that second-order methods are gaining interest from scientists because the field may be reaching a limit with the benefits first-order methods can provide. He believes the team’s findings show how second-order information can be extremely helpful if done in the right way.

New and better methods for optimizing neural networks have implications for a variety of fields, including economics, robotic control and in studies of climate, Kamzolov explains. “People have been solving these problems with simple but slow approaches, and now we have an opportunity to accelerate them,” he says. “And even a small acceleration can result in huge savings in training costs.”

Since more effective optimization methods can improve performance, scientists may be able to design smaller, more efficient networks that perform just as well as larger systems. For example, multimodal large language models that are built for text-to-image generation are based on neural networks with many layers. More efficient optimization techniques like the one developed by Agafonov, Kamzolov and their colleagues have the potential to reduce the need for layers in the network, decreasing the cost of training and running these systems.

Agafonov and Kamzolov said that they plan to continue to work on second-order methods and explore ways they can be applied to solve practical problems. “We can try to find some underlying ideas about the structure of neural nets where we can use second-order methods and the good approximations of curvature information in training for classification problems as well,” Agafonov says.