Recovery Guarantees for One-hidden-layer Neural Networks

In plain words

Suppose a hidden "teacher" network with one layer of k units produced the labels for your data. Can you train a student network of the same shape and get back the teacher's weights, not just a network that predicts similarly? This paper says yes, under conditions: the inputs are standard Gaussian random vectors, the labels are noiseless, and the input has at least as many dimensions as there are hidden units.

The trick has two stages. First, a cheap moment-based estimate, computed with tensor-style linear algebra, lands the student's weights in the right neighbourhood. Second, plain gradient descent, fed fresh data at each step, refines them. Near the truth the loss surface is bowl-shaped, so descent heads straight to the answer. The sample count grows roughly linearly with the input dimension, which is much cheaper than the cubic dependence of older tensor methods.

The guarantees are for a family of activation functions with three mild properties. Common choices such as ReLU, leaky ReLU, squared ReLU and sigmoid fit. The linear-speed convergence proof covers only the smooth ones.

TL;DR

For one-hidden-layer regression with Gaussian inputs and noiseless labels, the squared-loss Hessian is positive definite near the ground-truth weights for activations with three distilled properties §3, §4. For smooth homogeneous activations, tensor initialisation followed by gradient descent with resampling recovers the weights and output coefficients with sample size linear in the input dimension §5, §6. Theory only; the experiments are synthetic §7.

Problem

Inputs are x ~ N(0, I) in R^d. Labels come from a hidden one-layer network with k units §3:

y = \sum_{i=1}^{k} v_i^* \, \phi(w_i^{*\top} x) §3

The question is how many i.i.d. samples recover the parameters w_i* and v_i* up to permutation §1, §3. Training a one-hidden-layer net is NP-complete in general §3, so assumptions are required. The paper assumes Gaussian inputs, noiseless labels, k ≤ d so the first-layer vectors can be linearly independent §3, and an activation with the properties listed under Method §3.

Motivation: prior tensor methods need O(d³) samples §1, §2. Prior work by Tian (2017) handles the same setting but only in the population limit, with orthogonal first-layer weights and fixed second-layer weights §2. Kernel and over-parameterised approaches give generalisation or global optimality but not parameter recovery at this sample cost §1, §2.

Method

Activation properties §3. The activation φ is continuous, and at non-smooth points its derivative is the left derivative. Three properties are assumed:

1. Gradient bound. φ'(z) ≥ 0 and 0 ≤ φ'(z) ≤ L₁|z|^p for some L₁ > 0, p ≥ 0.
2. Expectation condition. Define α_q(σ) = E[φ'(σz) z^q] for q ∈ {0,1,2} and β_q(σ) = E[φ'²(σz) z^q] for q ∈ {0,2}, with z ~ N(0,1). Let ρ(σ) be the minimum of β₀ − α₀² − α₁², β₂ − α₁² − α₂², and α₀α₂ − α₁². Require ρ(σ) > 0 for all σ > 0.
3. Second derivative. Either (a) φ'' is bounded by L₂ (smooth case), or (b) φ'' is zero except at a finite number of points (piecewise linear case).

The paper states that ReLU, leaky ReLU, squared ReLU, and smooth non-decreasing activations with bounded symmetric derivative (sigmoid, tanh, erf) satisfy all three; linear and pure quadratic activations do not §3, app. C. Table 1 gives ρ(σ): exactly 0.091 for ReLU, 0.089 for leaky ReLU, 0.27σ for squared ReLU, and numerically 1.8e-4, 4.9e-2 and 5.1e-5 for sigmoid at σ = 0.1, 1 and 10 tab. 1. ρ is the constant that controls every later sample and radius bound.

Empirical risk and Hessian §4. The empirical risk is f̂_S(W) = (1/2|S|) Σ (Σ_i v_i* φ(w_i^T x) − y)², and the population risk replaces the average with an expectation §4. At W = W*, the off-diagonal Hessian blocks are E[v_j* v_l* φ'(w_jᵀx) φ'(w_lᵀx) x xᵀ], and the diagonal blocks, for piecewise-linear φ, are E[(v_j* φ'(w_j*ᵀx))² x xᵀ] §4. For non-smooth φ, φ'' is represented with Dirac functions §4.

Local strong convexity §4, app. D. For W within a radius of W*, with enough samples, the empirical Hessian satisfies, with probability 1 − d^(−Ω(t)) §4:

\Omega\!\left(\frac{v_{\min}^2\,\rho(\sigma_k)}{\kappa^2\lambda}\right) I \;\preceq\; \nabla^2 \hat f_S(W) \;\preceq\; O\!\left(k\, v_{\max}^2\, \sigma_1^{2p}\right) I §4

Definitions §4: σ_i are singular values of W* in decreasing order, κ = σ₁/σ_k, λ = (∏σ_i)/σ_k^k, v_max and v_min are the largest and smallest |v_i*|, ν = v_max/v_min, ρ = ρ(σ_k), and τ = (3σ₁/2)^{4p} / min over σ ∈ [σ_k/2, 3σ₁/2] of ρ²(σ). The formal radius and sample bound appear in the appendix app. D: the radius is v_min⁴ρ²(σ_k)/(k²κ⁵λ²v_max⁴σ₁^{4p}) times ‖W*‖, and the sample count is d·poly(log d, t)·k²v_max⁴τκ⁸λ²σ₁^{4p}/(v_min⁴ρ²(σ_k)). The informal version in the body writes both as unspecified poly(·) §4.

Linear convergence of gradient descent (smooth activations only, property 3a) §4, app. D. Gradient descent is run on the empirical risk with fresh samples at each step. For W^c inside the radius, the update W̃ = W^c − (1/M₀) ∇f̂_S(W^c), with m₀ = Θ(v_min²ρ(σ_k)/(κ²λ)) and M₀ = Θ(k v_max² σ₁^{2p}), satisfies with probability 1 − d^(−Ω(t)) §4, app. D:

\lVert \tilde W - W^*\rVert_F^2 \le \left(1 - \frac{m_0}{M_0}\right) \lVert W^c - W^*\rVert_F^2 §4

The proof covers the segment from W^c to W* by anchor points and a union bound app. D.

Tensor initialisation §5, alg. 1. The method uses moments of y and x:

• M₁ = E[y x]; M₂ = E[y(x⊗x − I)]; M₃ = E[y(x^⊗3 − x ⊗̃ I)]; M₄ = E[y(x^⊗4 − (x⊗x) ⊗̃ I + I ⊗̃ I)] §5, def. 1.
• With γ_j(σ) = E[φ(σz) z^j], the coefficients are m₁ = γ₁, m₂ = γ₂ − γ₀, m₃ = γ₃ − 3γ₁, m₄ = γ₄ + 3γ₀ − 6γ₂ §5, def. 1.
• Each M_j = Σ_i v_i* m_{j,i} w̄_i*^⊗j, where w̄ = w/‖w‖ §5, claim 1.

Assumption 1 (activation moments) §5: (1) if M_j ≠ 0 then m_{j,i} ≠ 0 for every i; (2) at least one of M₃, M₄ is non-zero; (3) if M₁ = M₃ = 0, φ is even; (4) if M₂ = M₄ = 0, φ is odd. The text asserts ReLU, leaky ReLU and squared ReLU satisfy this §5. For ReLU, M₃ = 0 §5. The value of m_{4,i} for ReLU appears only in commented-out source, not in the compiled text.

Procedure §5, alg. 1, app. E:

# Initialisation via tensor method  [alg. 1]
# Input:  samples S (linear in d), k
# Output: W0 (d, k), v0 (k,)
def initialise(samples):
    # Split into three disjoint subsets, each of size about d  [alg. 1] [app. E]
    subset_2, subset_3, subset_4 = partition(samples, 3)
    # Estimate the second-order moment P2 on subset_2  [§5, def. 2]
    p2_hat = empirical_p2(subset_2)            # (d, d)
    # Top-k eigen-subspace of p2_hat, via the power method  [alg. 3]
    basis = power_method(p2_hat, k)            # (d, k), orthonormal columns
    # Reduced third moment R3 = P3(V, V, V) on subset_3  [§5] [app. E]
    r3_hat = empirical_reduced_p3(subset_3, basis)   # (k, k, k)
    # Non-orthogonal tensor decomposition (KCL)  [alg. 1] [app. E, step b]
    u_hat = kcl_decompose(r3_hat)              # (k, k)
    # Recover norms, output signs and first-layer signs  [alg. 4] [app. E]
    w0, v0 = recover_magnitude_and_sign(basis, u_hat, subset_4)
    return w0, v0

Recovery step alg. 4 app. E: a random unit vector α is drawn. Two empirical moments Q̂₁ (from M_{l₁}(I, α, …, α)) and Q̂₂ (from M_{l₂}(V, V, α, …, α)) are formed on the two halves of subset_4. Least-squares fits give ẑ and r̂. Then v_i⁽⁰⁾ = sign(r̂_i c_{l₂}), s_i⁽⁰⁾ = sign(v_i⁽⁰⁾ ẑ_i c_{l₁}), and w_i⁽⁰⁾ = s_i⁽⁰⁾ (|ẑ_i / (c_{l₁} (αᵀV û_i)^{l₁−1})|)^{1/(p+1)} V û_i alg. 4. The choice of l₁, l₂ depends on which moments are non-zero, in three cases app. E.

Global algorithm §6, alg. 2. The sample set is split into T + 1 fresh subsets alg. 2. Initialisation runs on subset 0 alg. 2. The output weights v are then fixed at v⁽⁰⁾ for all later steps alg. 2.

# Globally converging algorithm  [alg. 2] [§6]
# Input:  samples S, k, number of steps T
# Output: W (d, k), v (k,)
def learn_one_hidden_layer(samples, num_steps):
    # Initialise from the first fresh subset  [alg. 2] [alg. 1]
    weights, v = initialise(subsets[0])
    # Use fresh samples at every step (resampling)  [alg. 2] [§4]
    for q in range(num_steps):
        batch = subsets[q + 1]
        grad = empirical_gradient(weights, v, batch)   # (d, k)
        # Step size eta <= 1 / (k * v_max**2 * sigma_1**(2p))  [§6]
        weights = weights - eta * grad
    return weights, v

For homogeneous activations the paper assumes v_i* ∈ {−1, 1}, using v φ(z) = (v/|v|) φ(|v|^{1/p} z) §4 (see Pitfalls on the exponent), so the output weights are recovered exactly once the estimate is close §5, app. E.

Hyperparameters

Results

Author-reported. Setup: synthetic Gaussian data, squared ReLU, k = 5, κ = 2, no resampling in the experiments §7. The paper's results are theorems rather than benchmark numbers. No real-data experiment is reported.

The figure values are shown only graphically; the text gives no numeric recovery rates fig. 1.

Paper vs code

No official code.

Searches: the paper, abstract and source contain no code link or project URL. GitHub API searches for "recovery guarantees one-hidden-layer", "tensor initialization gradient descent neural network zhong" and "one-hidden-layer neural network tensor" found no repository owned by an author. The experiments use the third-party code of Kuleshov et al. §7, which is not used here.

Pitfalls

• Sample count is stated three ways. The informal and formal theorems give d·poly(log d, t, …)/ε-type bounds §5, §6, app. D. The appendix sample discussion gives Õ(d) for P₂ and Õ(k³) for the reduced third moment app. E, while the proof of the tensor theorem uses Õ(d/ε²) for P₂ and Õ(k²/ε²) for R̂₃ app. E. An implementer should use the theorem's form and treat the exponents as unverified.
• Homogeneity exponent. The body says "p is the degree of homogeneity" and scales with v φ(z) = (v/|v|) φ(|v|^{1/p} z) §4, convex section. The appendix says the degree is p + 1, and the norm formula uses 1/(p+1) app. E, alg. 4. For squared ReLU (p = 1 by property 1) the degree is 2, so the exponent must be 1/(p+1); the body's 1/p gives the wrong identity.
• Experiments contradict the theory. The theory needs fresh samples at each step §4, §6 and step size ≤ 1/(k v_max² σ₁^{2p}) §6. The experiments use η = 0.02 and no resampling §7, and the paper says this "still works well" §7. Do not infer that resampling can be dropped.
• Non-smooth activations. Linear convergence is proven only for smooth activations (property 3a) §4, app. D. ReLU is covered by the tensor initialisation and the local Hessian bound, but not by the linear-rate theorem §6.
• Moment selection is not specified. The algorithm picks l₁, l₂ as the first non-zero M_j alg. 4. In practice M_j = 0 must be decided from samples, and the paper does not say how. For finite samples the choice is unverified.
• Constants are hidden. The Hessian radius, sample bounds, and step constant are written as Θ(·) or poly(·) §4, app. D. The exponents of k, κ, λ, ν and ρ are not stated in the informal forms.
• ρ controls the cost. Sample count and radius scale with 1/ρ(σ_k)² app. D. Table 1 gives ρ = 5.1e-5 for sigmoid at σ = 10 tab. 1, so the bounds become very large for that case. The paper gives no closed form for sigmoid or tanh.
• Noise and generality. Labels are noiseless and inputs are standard Gaussian. The paper says a general covariance is handled by estimating and whitening it, "with some loss of accuracy", without a bound §3. Classification is reduced to regression only with a good teacher §3. The results do not cover noisy labels.
• Linear independence. The analysis needs k ≤ d and linearly independent first-layer vectors. Rank-deficient W* gives λ → ∞ and κ → ∞, and the positive-definite property fails §4.
• Random α. The recovery step uses a random unit vector α and needs αᵀ w̄_i* ≠ 0 §5, app. E. The paper says this holds with probability one but does not state how to check it.
• Figure 1(a)–(b) are qualitative. Numbers are not given in the text.

Lineage

• Builds on the tensor method of Janzamin, Sedghi and Anandkumar (2015) for one-hidden-layer nets, whose sample complexity is O(d³/ε²) §2. Builds on Zhong, Jain and Dhillon for the power-method subspace step §5. Uses the non-orthogonal tensor decomposition of Kuleshov et al. (2015) §5.
• Compared against: Tian, which considers the same Gaussian teacher setting but orthogonal first-layer weights, fixed second-layer weights and population risk only §2. The paper's own §1 comparisons are the only ones given.
• Superseded by: not stated.