A Unified Framework for Lifted Training and Inversion Approaches

2025-10-10

AI Summary

To address inherent limitations of backpropagation—including vanishing/exploding gradients, poor handling of nonsmooth activations, and limited parallelizability—this paper proposes a unified training framework based on constrained optimization. It formulates deep neural network training as a high-dimensional constrained optimization problem, relaxes constraints via penalty terms, and improves manifold conditioning using Bregman distances. The framework uniquely unifies proximal-type nonsmooth activations, distributed optimization, and joint forward-inverse problem solving, integrating auxiliary coordinate methods, Fenchel duality, and block-coordinate descent, while supporting acceleration and adaptive stochastic/deterministic updates. Experimental validation on standard imaging tasks demonstrates that, compared to conventional backpropagation, the method achieves significantly improved convergence speed, numerical stability, and robustness—particularly for networks employing proximal activations.

Technology Category

• Search and Optimization: Non-convex Optimization
• Computer Vision: Learning & Optimization for CV
• Machine Learning: Optimization

Application Category

• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Social Networks and Social Media: Influence propagation, information diffusion, and the prediction on networks

Abstract

The training of deep neural networks predominantly relies on a combination of gradient-based optimisation and back-propagation for the computation of the gradient. While incredibly successful, this approach faces challenges such as vanishing or exploding gradients, difficulties with non-smooth activations, and an inherently sequential structure that limits parallelisation. Lifted training methods offer an alternative by reformulating the nested optimisation problem into a higher-dimensional, constrained optimisation problem where the constraints are no longer enforced directly but penalised with penalty terms. This chapter introduces a unified framework that encapsulates various lifted training strategies, including the Method of Auxiliary Coordinates, Fenchel Lifted Networks, and Lifted Bregman Training, and demonstrates how diverse architectures, such as Multi-Layer Perceptrons, Residual Neural Networks, and Proximal Neural Networks fit within this structure. By leveraging tools from convex optimisation, particularly Bregman distances, the framework facilitates distributed optimisation, accommodates non-differentiable proximal activations, and can improve the conditioning of the training landscape. We discuss the implementation of these methods using block-coordinate descent strategies, including deterministic implementations enhanced by accelerated and adaptive optimisation techniques, as well as implicit stochastic gradient methods. Furthermore, we explore the application of this framework to inverse problems, detailing methodologies for both the training of specialised networks (e.g., unrolled architectures) and the stable inversion of pre-trained networks. Numerical results on standard imaging tasks validate the effectiveness and stability of the lifted Bregman approach compared to conventional training, particularly for architectures employing proximal activations.

Problem

• Addresses vanishing and exploding gradients in deep networks
• Enables distributed optimization for neural network training
• Handles non-differentiable activations in inverse problem solving

Innovation

• Unified framework combines various lifted training strategies
• Uses Bregman distances for distributed optimization and conditioning
• Applies block-coordinate descent for training and inversion problems