Global Convergence and Rich Feature Learning in $L$-Layer Infinite-Width Neural Networks under $mu$P Parametrization

2025-03-12

AI Summary

研究解决深度神经网络在无限宽度下如何实现全局收敛和丰富特征学习的问题，采用张量程序框架和最大更新参数化方法，通过随机梯度下降训练，确保特征线性独立并捕捉数据信息。

Technology Category

• Machine Learning: Deep Learning Algorithms
• Computer Vision: Learning & Optimization for CV
• Natural Language Processing: Learning & Optimization for NLP

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
• Web Mining and Content Analysis: Large pretrained models with web data

Abstract

Despite deep neural networks' powerful representation learning capabilities, theoretical understanding of how networks can simultaneously achieve meaningful feature learning and global convergence remains elusive. Existing approaches like the neural tangent kernel (NTK) are limited because features stay close to their initialization in this parametrization, leaving open questions about feature properties during substantial evolution. In this paper, we investigate the training dynamics of infinitely wide, $L$-layer neural networks using the tensor program (TP) framework. Specifically, we show that, when trained with stochastic gradient descent (SGD) under the Maximal Update parametrization ($mu$P) and mild conditions on the activation function, SGD enables these networks to learn linearly independent features that substantially deviate from their initial values. This rich feature space captures relevant data information and ensures that any convergent point of the training process is a global minimum. Our analysis leverages both the interactions among features across layers and the properties of Gaussian random variables, providing new insights into deep representation learning. We further validate our theoretical findings through experiments on real-world datasets.

Problem

• Understanding feature learning and global convergence in deep networks
• Exploring training dynamics of infinite-width neural networks
• Validating theoretical insights with real-world dataset experiments

Innovation

• Uses Maximal Update parametrization (μP) for training
• Leverages tensor program framework for analysis
• Ensures global convergence with rich feature learning