An alternative approach to train neural networks using monotone variational inequality

TL;DR

This paper introduces a monotone variational inequality (VI) approach to neural network training, replacing traditional gradient-based updates with a monotone operator-derived vector field (SVI). It provides convergence guarantees for last-layer fine-tuning when the VI is strongly monotone and demonstrates feasible extensions to training multiple layers via a backward unrolling scheme. Theoretical results cover both strongly monotone (κ>0) and merely monotone (κ=0) cases, including conditions under which the operator reduces to the gradient. Empirical results across synthetic data, graph neural networks, and image datasets show that |SVI| often achieves faster early convergence and competitive final performance compared to SGD, with notable benefits for fine-tuning pre-trained models and deploying large architectures.

Abstract

We propose an alternative approach to neural network training using the monotone vector field, an idea inspired by the seminal work of Juditsky and Nemirovski [Juditsky & Nemirovsky, 2019] developed originally to solve parameter estimation problems for generalized linear models (GLM) by reducing the original non-convex problem to a convex problem of solving a monotone variational inequality (VI). Our approach leads to computationally efficient procedures that converge fast and offer guarantee in some special cases, such as training a single-layer neural network or fine-tuning the last layer of the pre-trained model. Our approach can be used for more efficient fine-tuning of a pre-trained model while freezing the bottom layers, an essential step for deploying many machine learning models such as large language models (LLM). We demonstrate its applicability in training fully-connected (FC) neural networks, graph neural networks (GNN), and convolutional neural networks (CNN) and show the competitive or better performance of our approach compared to stochastic gradient descent methods on both synthetic and real network data prediction tasks regarding various performance metrics.

Figures (21)

• Figure 1: Diagram to show fine-tuning by training the last two layers for a neural network. Black lines indicate frozen parameters and dark green lines indicate parameters to be fine-tuned by training; in this plot, the last two layers are to be trained.
• Figure 2: Gradient descent (GD) vs. SVI for fine-tuning an arbitrary layer $l$, $L'\leq l\leq L$, via backward unrolling scheme: the difference lies in the skipping of differentiation of the activation $\phi$ with respect to pre-activation values. The notation $\circ$ denotes function composition.
• Figure 3: Losses under MSE objective
• Figure 4: Losses under NLL objective
• Figure 5: Accuracies on the same test data

Theorems & Definitions (14)

• Lemma 4.1
• Lemma 4.2: Parameter recovery guarantee
• Theorem 4.3: Prediction error using strongly monotone $F$
• Remark 4.4: When does $\kappa>0$?
• Theorem 4.5: Prediction error using monotone $F$
• Proposition 4.6: The equivalence between SVI and parameter gradient
• Remark 5.1: Effect on training dynamics
• Remark 5.2: Implementation caveats
• Remark 5.3: Extension to GNN
• proof : Proof of Lemma \ref{['lem1']}