Efficient and Flexible Neural Network Training through Layer-wise Feedback Propagation

TL;DR

Layer-wise Feedback Propagation (LFP) is a novel training method that uses explainability techniques to assign neuron-specific rewards, enabling efficient, flexible, and non-gradient-based neural network training with theoretical convergence guarantees.

Contribution

Introduces LFP, a new training principle that decomposes rewards to neurons for efficient, flexible, and gradient-free neural network optimization, applicable to various architectures.

Findings

LFP achieves comparable computational complexity to gradient descent.

LFP produces sparse, memory- and energy-efficient parameter updates.

LFP demonstrates effectiveness in neural network pruning and training Spiking Neural Networks.

Abstract

Gradient-based optimization has been a cornerstone of machine learning that enabled the vast advances of Artificial Intelligence (AI) development over the past decades. However, this type of optimization requires differentiation, and with recent evidence of the benefits of non-differentiable (e.g. neuromorphic) architectures over classical models w.r.t. efficiency, such constraints can become limiting in the future. We present Layer-wise Feedback Propagation (LFP), a novel training principle for neural network-like predictors that utilizes methods from the domain of explainability to decompose a reward to individual neurons based on their respective contributions. Leveraging these neuron-wise rewards, our method then implements a greedy approach reinforcing helpful parts of the network and weakening harmful ones. While having comparable computational complexity to gradient descent, LFP does not require gradient computation and generates sparse and thereby memory- and energy-efficient parameter updates and models. We establish the convergence of LFP theoretically and empirically, demonstrating its effectiveness on various models and datasets. Via two applications - neural network pruning and the approximation-free training of Spiking Neural Networks (SNNs) - we demonstrate that LFP combines increased efficiency in terms of computation and representation with flexibility w.r.t. choice of model architecture and objective function. Our code is available at https://github.com/leanderweber/layerwise-feedback-propagation.