We investigate the convergence guarantee of two-layer neural network training with Gaussian randomly masked inputs. This scenario corresponds to Gaussian dropout at the input level, or noisy input training common in sensor networks, privacy-preserving training, and federated learning, where each user may have access to partial or corrupted features. Using a Neural Tangent Kernel (NTK) analysis, we demonstrate that training a two-layer ReLU network with Gaussian randomly masked inputs achieves linear convergence up to an error region proportional to the mask's variance. A key technical contribution is resolving the randomness within the non-linear activation, a problem of independent interest.
Figures & tables
Figure 1: (a). Effect of the noise standard deviation κ on the shape of the smoothed activation function σ^(z;κ)=z⋅Φ1(z/(κ∥w⊙x∥2)) , where z=w⊤x . For this visualization, ∥w⊙x∥2 is held constant at 1.0 . As κ increases, the activation becomes progressively smoother compared to the standard ReLU (dotted black line). For small κ (e.g., κ=0.01 ), σ^ closely approximates the standard ReLU. (b). Theoretical smoothed activation σ^(w,x) versus its empirical estimate Ec[σ(w⊤(x⊙c))] for a fixed pre-activation value w⊤x≈0.77 (actual value depends on fixed w,x ) as the noise standard deviation κ varies. The close match across a range of—relatively small— κ values validates the theoretical model for σ^ . Note that this behavior consistently follows empirically for different w,x values.
Figure 2: Smoothed ReLU under multiplicative Gaussian input masking for fixed κ=0.2 , where z=w⊤x and σ=κ∥w⊙x∥2 . (a) Exact closed-form expectation σ~(w,x)=Ec[σ(w⊤(x⊙c))]=zΦ(z/σ)+σφ(z/σ) (as shown in 19 )) matches the Monte Carlo estimate. (b) Proxy smoothed activation σ^(w,x)=zΦ(z/σ) (used in Theorem 4.2 ) differs mainly near z≈0 due to the missing σφ(z/σ) term.
Figure 3: (a). Training loss L(Wk) (log-scale) versus training iteration for a two-layer ReLU network ( n=500,d=20,m=100 ) trained with full-batch gradient descent under different levels of input multiplicative Gaussian noise standard deviation κ . (b). Distributed training with Gaussian mask for differen κ and number of local steps.
Figure 4: Test accuracy versus multiplicative Gaussian noise strength κ for (a) a 1-hidden-layer MLP and (b) a CNN, trained on CIFAR-10. Small noise levels ( κ≈0.2) can improve generalization for the MLP, likely due to regularization effects. In contrast, the CNN exhibits robustness by maintaining baseline accuracy. Beyond this point, accuracy degrades monotonically for both architectures as noise corrupts the training signal.
Figure 5: Attack AUC on the target model when it is (a) an MLP and (b) a CNN. Higher values indicate greater privacy leakage. Training with MG noise (larger κ ) consistently reduces attack success.
Figure 6: MLP target. Multiplicative Gaussian noise provides resilience against the attacks as κ increases.
Figure 7: CNN target. We observe a similar trend as the MLP
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 8: ℓ2 -norms of the gradient components (left) and residual error bound check (right).
Figure 9: Privacy-utility tradeoff for the Standard CNN. The left panel shows Attack Precision vs. Accuracy, and the right panel shows Attack Recall vs. Accuracy.
Figure 10: Privacy-utility tradeoff for the Improved CNN. Unlike the Standard CNN, the MG noise curve stays much closer to the DP-SGD curve across the accuracy range
In this paper, we study the gradient descent dynamics for jointly training both layers of a one-hidden-layer ReLU network to fit a linear target function. Concretely, we consider a realizable setting where inputs are drawn i.i.d. from a Gaussian distribution and labels follow a planted linear model. This stylized framework captures salient features of end-to-end training in inverse problems and certain auto-encoder models. Despite its apparent simplicity, the dynamics remain poorly understood, in part because the loss landscape contains multiple non-strict saddle points, making it unclear why gradient descent from random initialization reliably escapes bad stationary regions. We provide a detailed characterization of the optimization landscape and prove that gradient descent from a moderately small random initialization-simultaneously training both layers-converges to a global minimizer at a linear rate with order-wise optimal sample complexity. Our analysis tracks the trajectory through three phases: an alignment phase in which hidden weights progressively align with the planted direction while the output weights maintain the correct sign pattern; a growth phase in which the norms of both layers increase while preserving alignment; and a local refinement phase in which the aligned neurons rapidly converge to the planted direction, yielding fast local convergence. To rigorously show that GD avoids non-strict saddles, we develop trajectory-level control arguments for the end-to-end dynamics. In addition, we establish novel uniform concentration results that hold along the entire trajectory, and are essential for obtaining order-wise optimal sample complexity. We corroborate our theory with extensive experiments across a range of configurations.
Berk Tinaz, Changzhi Xie, Mahdi Soltanolkotabi
Department of Electrical and Computer Engineering University of Southern California Los Angeles, CA, USA · Department of Computer Science University of Southern California Los Angeles, CA, USA
We establish the first population risk bounds for Kolmogorov-Arnold Networks (KANs) trained by mini-batch SGD with gradient clipping, covering non-private SGD as well as differentially private SGD (DP-SGD) with Gaussian perturbations that interpolate between independent and temporally correlated noise. This setting is substantially closer to practice than prior KAN theory along two axes: training is by mini-batch SGD, the standard recipe for modern networks, rather than full-batch gradient descent (GD); and correlated-noise mechanisms have empirically shown a more favorable privacy-utility tradeoff than independent-noise mechanisms. Our results cover the corresponding full-batch GD and independent-noise DP-GD results for KANs by Wang et al. (2026), while yielding sharper fixed-second-layer specializations. The technical core is a new analysis route for correlated-noise DP training in the non-convex regime. Temporal dependence breaks the conditional-centering structure underlying standard one-step SGD arguments, and the projection step obstructs the exact cancellation structure of correlated perturbations. We address these difficulties through an auxiliary unprojected dynamics, a shifted iterate that absorbs the current noise perturbation, and a high-probability bootstrap certifying projection inactivity. Combining this optimization analysis with a stability-based generalization argument yields the stated population risk bounds. To the best of our knowledge, this is the first optimization and population risk analysis of a correlated-noise mechanism for DP training beyond convex learning, in particular for neural networks.
Puyu Wang, Jan Schuchardt, Nikita Kalinin +4
RPTU Kaiserslautern-Landau, Kaiserslautern, Germany · Machine Learning Research, Morgan Stanley · Institute of Science and Technology, Klosterneuburg, Austria +1
We prove global exponential (linear) convergence with an explicit rate in the rich scaling for wide two-layer linear networks trained with smooth Polyak-Lojasiewicz predictor losses. Gradient flow in the factors closes exactly in terms of a finite-dimensional Bures flow of the neuron law covariance, in which the predictor dynamics are preconditioned by hidden covariance blocks. Mean-field conservation laws provide uniform spectral lower bounds on the hidden preconditioning blocks when the initial covariance satisfies a spectral support gap condition. This condition encompasses positive definiteness while still allowing for singular initializations. For an initial covariance Σ0=σ2Id, the loss converges to the global minimum with linear rate at least 4σ2κ, where κ is the PL constant. We establish stability of this rate under finite-width sampling, as well as global convergence of factor gradient descent for an explicit stepsize interval depending on smoothness, the initial loss, and conserved spectral margins. Our argument extends layerwise to deep linear ResNets, subject to a residual-path bound. In the case of heavy-ball momentum, training dynamics close instead over positions and velocities in terms of a lifted phase covariance. Linear convergence holds under an explicit condition on the energy and damping, specifying a window of admissible dampings. For two-scale white initializations, this interval is nonempty for sufficiently large position scales, with a fixed initial loss gap and velocity covariance. Numerical experiments illustrate the covariance geometry and compare the predicted and observed rates.