A generalized neural tangent kernel for surrogate gradient learning

2024-05-24

Neural Information Processing Systems

AI Summary

To address the failure of standard gradient-based training caused by non-differentiable activation functions (e.g., binary or spiking activations), this paper introduces the Surrogate Gradient Neural Tangent Kernel (SG-NTK), the first rigorous extension of NTK theory to training dynamics involving discontinuous activations and surrogate derivatives. Methodologically, leveraging functional analysis and kernel methods, we establish the existence and convergence of SG-NTK in the infinite-width limit and empirically validate its dynamical characterization capability on finite-width networks. Theoretically, we prove that SG-NTK exactly captures the surrogate gradient learning process—resolving the long-standing lack of theoretical foundation for surrogate gradient learning (SGL). Experimentally, on sign-activated networks, predictions derived from SG-NTK closely match those of kernel regression, confirming both theoretical soundness and practical applicability.

Technology Category

• Machine Learning: Kernel Methods
• Cognitive Modeling & Cognitive Systems: Neural Spike Coding
• Natural Language Processing: Learning & Optimization for NLP

Application Category

• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Graph neural networks and deep learning approaches for Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

State-of-the-art neural network training methods depend on the gradient of the network function. Therefore, they cannot be applied to networks whose activation functions do not have useful derivatives, such as binary and discrete-time spiking neural networks. To overcome this problem, the activation function's derivative is commonly substituted with a surrogate derivative, giving rise to surrogate gradient learning (SGL). This method works well in practice but lacks theoretical foundation. The neural tangent kernel (NTK) has proven successful in the analysis of gradient descent. Here, we provide a generalization of the NTK, which we call the surrogate gradient NTK, that enables the analysis of SGL. First, we study a naive extension of the NTK to activation functions with jumps, demonstrating that gradient descent for such activation functions is also ill-posed in the infinite-width limit. To address this problem, we generalize the NTK to gradient descent with surrogate derivatives, i.e., SGL. We carefully define this generalization and expand the existing key theorems on the NTK with mathematical rigor. Further, we illustrate our findings with numerical experiments. Finally, we numerically compare SGL in networks with sign activation function and finite width to kernel regression with the surrogate gradient NTK; the results confirm that the surrogate gradient NTK provides a good characterization of SGL.

Problem

• Analyzes surrogate gradient learning for non-differentiable activation functions
• Generalizes neural tangent kernel to surrogate gradient descent
• Validates surrogate gradient NTK with numerical experiments

Innovation

• Generalized neural tangent kernel for surrogate gradients
• Surrogate gradient NTK enables SGL analysis
• Mathematically rigorous NTK generalization for SGL