On Learnability via Gradient Method for Two-Layer ReLU Neural Networks in Teacher-Student Setting

Theoretical understanding of deep learning remains incomplete, especially regarding optimization landscapes and generalization. Early works like Li et al. (2018) and Zhang et al. (2019) explored convergence and over-parameterization effects, but lacked comprehensive guarantees for finite-width, non-smooth activations. NTK (Jacot et al., 2018) provided a linearized view, yet failed to explain feature learning. Mean field analysis (Mei et al., 2019) considered infinite-width limits, offering insights but limited applicability. Teacher-student models have been used to analyze parameter recovery, but often under restrictive assumptions. This paper advances the field by combining measure representation with sparse regularization, addressing the challenge of non-smooth ReLU activations and finite network widths.

The core challenge is to establish whether gradient descent can reliably recover the true teacher network parameters in finite-width, non-smooth ReLU networks. Existing theories often rely on idealized assumptions like infinite width or smooth activations, which do not reflect practical scenarios. The non-convex landscape, coupled with the non-differentiability of ReLU, complicates convergence analysis. Moreover, understanding the conditions under which the learned network approximates the teacher with high accuracy remains unresolved, especially when the teacher parameters are not orthogonal or data is noisy.

This work introduces a measure space reformulation, transforming the neural network training into a sparse measure optimization problem. It applies BLASSO theory to analyze the sparsity and recovery guarantees, leveraging dual certificates to establish global convergence. The approach explicitly relates regularization strength to parameter estimation error, providing a rigorous guarantee that the gradient descent can recover the teacher parameters accurately. Unlike prior work limited to infinite-width or smooth activations, this method handles finite-width, non-smooth ReLU networks, bridging a crucial gap in theory.

• �� Reformulate the two-layer ReLU network as a measure ν on the sphere Sd−1, representing the network as an integral over this measure. • Define a regularized empirical risk minimization problem on the measure space, incorporating total variation regularization to promote sparsity. • Use Fenchel duality to derive the dual problem, constructing a dual certificate function that verifies optimality conditions. • Prove that under sufficient over-parameterization and small regularization λ, the optimal measure is a sparse sum of Dirac measures close to teacher parameters. • Design a norm-dependent gradient descent algorithm, updating parameters with step sizes scaled by their norms, ensuring convergence in measure space. • Validate the theoretical results through synthetic experiments, demonstrating parameter recovery and network size reduction.

Synthetic datasets generated from teacher networks with widths ranging from 5 to 20. Training involved varying regularization λ and sample size n, observing the parameter error and output approximation. The experiments compared the convergence behavior of the measure-based gradient descent with baseline methods, confirming the theoretical predictions. Metrics included parameter distance, output error, and sparsity level. Results showed that decreasing λ improved recovery accuracy, with the learned network parameters approaching the teacher’s, and the network size shrinking to the teacher’s width, validating the measure representation’s effectiveness.

The experiments confirmed that for sufficiently large n, the optimal solutions of the regularized problem approximate teacher parameters within O(mλ^2). As λ approaches zero, the network accurately recovers the teacher’s weights, with the network size reducing to the true width m. Gradient descent in measure space demonstrated global convergence, even with non-smooth ReLU activation, and maintained robustness across different data distributions. These results support the theoretical claims, illustrating the potential of measure-based regularization for neural network identifiability.

Immediate applications include neural network parameter recovery, model compression, and transfer learning, especially in settings requiring precise parameter estimation from limited data. The theoretical framework also informs the design of regularization strategies for better generalization and interpretability. Long-term, the approach could facilitate the development of more robust, explainable deep models, and guide the design of training algorithms that guarantee convergence to meaningful solutions in complex, real-world tasks.

The analysis assumes teacher parameters are orthogonal, which simplifies the dual certificate construction but may not hold in practice. Tuning the regularization parameter λ remains challenging in real applications. The current focus on shallow, two-layer networks limits direct extension to deep architectures. Future work must address these limitations by relaxing assumptions, incorporating noise, and scaling algorithms for larger models.