Accelerating Training with Neuron Interaction and Nowcasting Networks

2024-09-06

arXiv.org

AI Summary

This work addresses the high computational cost and training instability associated with learnable update rules in neural network optimization. To this end, we propose NiNo—a novel architecture that integrates a graph neural network (GNN)-based parameter state nowcasting mechanism into the Adam optimization framework. NiNo explicitly models structural dependencies among neurons by leveraging the network’s inherent topology, thereby overcoming the limitation of conventional weight prediction methods that ignore graph-structured relationships. Notably, it is the first approach to unify neuron interaction modeling with GNNs, making it especially suitable for complex architectures such as Transformers. Experiments across vision and language tasks demonstrate that NiNo accelerates Adam-based training by up to 50%, while preserving optimization stability and model generalization performance.

Technology Category

• Natural Language Processing: Learning & Optimization for NLP
• Search and Optimization: Learning to Search
• Machine Learning: Optimization

Application Category

• 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
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Neural network training can be accelerated when a learnable update rule is used in lieu of classic adaptive optimizers (e.g. Adam). However, learnable update rules can be costly and unstable to train and use. Recently, Jang et al. (2023) proposed a simpler approach to accelerate training based on weight nowcaster networks (WNNs). In their approach, Adam is used for most of the optimization steps and periodically, only every few steps, a WNN nowcasts (predicts near future) parameters. We improve WNNs by proposing neuron interaction and nowcasting (NiNo) networks. In contrast to WNNs, NiNo leverages neuron connectivity and graph neural networks to more accurately nowcast parameters. We further show that in some networks, such as Transformers, modeling neuron connectivity accurately is challenging. We address this and other limitations, which allows NiNo to accelerate Adam training by up to 50% in vision and language tasks.

Problem

• Accelerate neural network training using learnable update rules.
• Improve weight nowcasting with neuron interaction and graph networks.
• Enhance training speed by 50% in vision and language tasks.

Innovation

• Uses weight nowcaster networks for parameter prediction
• Incorporates neuron interaction via graph neural networks
• Accelerates Adam training by up to 50%