Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

A Derandomization Framework for Structure Discovery: Applications in Neural Networks and Beyond

Nikos Tsikouras, Yorgos Pantis, Ioannis Mitliagkas, Christos Tzamos

TL;DR

This work addresses why neural networks trained with SGD exhibit low-rank structure in the first-layer weights. It introduces a general derandomization lemma applied to a broad class of smooth objectives with Gaussian inputs, showing that after reaching a rho-SOSP, the first-layer weights become effectively low-rank. The framework extends to arbitrary network depth with trainable biases and tiny regularization, and yields deterministic derandomizations for MAXCUT and Johnson-Lindenstrauss embeddings. Empirical results on two-layer networks corroborate the theory by showing convergence to the principal subspace, highlighting the framework's impact on understanding implicit regularization and feature learning.

Abstract

Understanding the dynamics of feature learning in neural networks (NNs) remains a significant challenge. The work of (Mousavi-Hosseini et al., 2023) analyzes a multiple index teacher-student setting and shows that a two-layer student attains a low-rank structure in its first-layer weights when trained with stochastic gradient descent (SGD) and a strong regularizer. This structural property is known to reduce sample complexity of generalization. Indeed, in a second step, the same authors establish algorithm-specific learning guarantees under additional assumptions. In this paper, we focus exclusively on the structure discovery aspect and study it under weaker assumptions, more specifically: we allow (a) NNs of arbitrary size and depth, (b) with all parameters trainable, (c) under any smooth loss function, (d) tiny regularization, and (e) trained by any method that attains a second-order stationary point (SOSP), e.g.\ perturbed gradient descent (PGD). At the core of our approach is a key $\textit{derandomization}$ lemma, which states that optimizing the function $\mathbb{E}_{\mathbf{x}} \left[g_θ(\mathbf{W}\mathbf{x} + \mathbf{b})\right]$ converges to a point where $\mathbf{W} = \mathbf{0}$, under mild conditions. The fundamental nature of this lemma directly explains structure discovery and has immediate applications in other domains including an end-to-end approximation for MAXCUT, and computing Johnson-Lindenstrauss embeddings.

A Derandomization Framework for Structure Discovery: Applications in Neural Networks and Beyond

TL;DR

This work addresses why neural networks trained with SGD exhibit low-rank structure in the first-layer weights. It introduces a general derandomization lemma applied to a broad class of smooth objectives with Gaussian inputs, showing that after reaching a rho-SOSP, the first-layer weights become effectively low-rank. The framework extends to arbitrary network depth with trainable biases and tiny regularization, and yields deterministic derandomizations for MAXCUT and Johnson-Lindenstrauss embeddings. Empirical results on two-layer networks corroborate the theory by showing convergence to the principal subspace, highlighting the framework's impact on understanding implicit regularization and feature learning.

Abstract

Understanding the dynamics of feature learning in neural networks (NNs) remains a significant challenge. The work of (Mousavi-Hosseini et al., 2023) analyzes a multiple index teacher-student setting and shows that a two-layer student attains a low-rank structure in its first-layer weights when trained with stochastic gradient descent (SGD) and a strong regularizer. This structural property is known to reduce sample complexity of generalization. Indeed, in a second step, the same authors establish algorithm-specific learning guarantees under additional assumptions. In this paper, we focus exclusively on the structure discovery aspect and study it under weaker assumptions, more specifically: we allow (a) NNs of arbitrary size and depth, (b) with all parameters trainable, (c) under any smooth loss function, (d) tiny regularization, and (e) trained by any method that attains a second-order stationary point (SOSP), e.g.\ perturbed gradient descent (PGD). At the core of our approach is a key lemma, which states that optimizing the function converges to a point where , under mild conditions. The fundamental nature of this lemma directly explains structure discovery and has immediate applications in other domains including an end-to-end approximation for MAXCUT, and computing Johnson-Lindenstrauss embeddings.
Paper Structure (55 sections, 8 theorems, 82 equations, 7 figures, 2 algorithms)

This paper contains 55 sections, 8 theorems, 82 equations, 7 figures, 2 algorithms.

Table of Contents

  1. Introduction
  2. Importance of training the biases.
  3. Discussion on $\rho$-SOSPs.
  4. Summary of our contributions.
  5. Informal version of Lemma \ref{['lem:main lemma']}
  6. Structure discovery in NNs.
  7. Our results in other domains.
  8. Notation and Preliminaries
  9. Notation.
  10. Main Contribution: Key Derandomization Lemma
  11. Proof sketch:
  12. Structure Discovery in Neural Networks
  13. Discovering Structure in Neural Networks via SOSPs
  14. Proof.
  15. Proof.
  16. ...and 40 more sections

Key Result

Lemma 3.1

Let $\bm{x} \sim \mathcal{N}(\bm{0}, \bm{I}_d)$ be a standard multivariate Gaussian random variable. For the objective function defined in Equation eq:regularized objective, with $\lambda > \frac{\sqrt{K\rho}}{2}$ where $g_{\theta}(\cdot)$ satisfies Assumption assump:1, any $\rho$-SOSP satisfies $\|

Figures (7)

  • Figure 1: Plot of the global minimizer of $f(w,0) = \mathbb{E}[(\text{ReLU}^3(wx) - 1)^2] + \lambda w^2$ as a function of the regularization parameter $\lambda$.
  • Figure 2: Optimal $w$ vs. $\lambda$ for $f(w, b) = \mathbb{E}[(\text{ReLU}^3(wx+b) - 1)^2] + \lambda w^2$ when biases are trained. The solid line shows convergence to 0 at an exact SOSP.
  • Figure 3: Two-layer $\text{ReLU}_{2}$ network of width $h=1000$ and $d = 2$ for the task of recovering a $\tanh$ single-index teacher model. We observe convergence of weights $\bm{W}$ to the principal subspace.
  • Figure 4: Progress of the cut value over iterations.
  • Figure 5: Evolution of the maximum $\sigma^2$ value over iterations.
  • ...and 2 more figures

Theorems & Definitions (15)

  • Definition 2.1
  • Definition 2.2: Approximate second-order stationarity
  • Lemma 3.1: Key Derandomization Lemma
  • Remark 3.2
  • Theorem 4.1
  • Theorem 4.2
  • Theorem 4.3
  • Remark 5.1
  • Theorem 5.2: Derandomized Approximation for MAXCUT
  • Definition 5.3: JL guarantee
  • ...and 5 more