Shallow ReLU Networks Get Sharper Generalization Bounds

Analysis by the aitrendblend editorial team · JMLR 2026 · Optimization & Learning Theory

A researcher increases the hidden width of a small classifier because the proof seems to demand it. Training already works with far fewer neurons. The gap raises an awkward question. Does the network need all those parameters, or does the analysis need them?

Optimization and Generalization of Gradient Descent for Shallow ReLU Networks with Minimal Width addresses that gap without replacing ReLU by a smooth activation. Yunwen Lei, Puyu Wang, Yiming Ying and Ding-Xuan Zhou combine optimization analysis with a more selective measure of model complexity.

The publication appeared in the Journal of Machine Learning Research in January 2026. The authors’ affiliations include the University of Hong Kong, RPTU Kaiserslautern-Landau and the University of Sydney. Their result belongs in our Optimization & Learning Theory collection, where the useful question is what a guarantee actually covers.

The headline is appealing. A shallow ReLU network can achieve a population risk bound that scales almost inversely with the number of training examples, under a separability assumption and suitable width. The important part is how that statement is earned. It depends on the training trajectory, the geometry of the tangent features and the difference between average risk and a particular trained checkpoint.

Training loss is only half the explanation

A network can fit its training examples and still predict poorly on fresh data. Optimization asks whether the training procedure finds a model with low empirical loss. Generalization asks how the resulting predictions relate to the underlying data distribution. A theory that answers only the first question leaves the main statistical question open.

The source makes the distinction explicit by decomposing population risk into empirical risk and a generalization gap. These are bookkeeping terms, but the decomposition gives the proof a clear job. Control the training behavior, then control the difference between the sample and the population for the models that training can reach.

Equation 1. Population risk, empirical risk and their decomposition, following Section 3.

Neither quantity is simply classification accuracy. For the main specialization, the loss is logistic. It scores the signed output continuously, distinguishing a confident correct decision from a weakly correct one. A bound on this loss contains useful classification information, but it should not be presented as a reported accuracy percentage.

This distinction is easy to lose when a paper says that its bound improves by a factor involving the sample size. That is a comparison of mathematical rates. It is not a claim that the same network becomes that much more accurate in an experiment. No measured accuracy curve is supplied here.

The network is small in depth and specific in design

The model has one hidden layer. Its output is a signed sum of ReLU activations, with the hidden weights trained and the output weights fixed. There are no biases in the stated architecture. Labels are binary, taking the values negative one and positive one, and inputs have unit Euclidean norm.

Equation 2. The shallow network in Section 3. Only the hidden parameter matrix represented by (W) is updated.

The initialization uses paired hidden neurons. Each pair begins with the same Gaussian hidden weight and opposite output signs. Their contributions cancel, giving zero initial network output. The cancellation does not prevent learning because the hidden weights receive updates with opposite signs.

The description of initialization contains an abbreviated sign notation after specifying the scaled output coefficients. The implementation below retains the explicitly stated magnitude of one over the square root of width. Otherwise, the gradient size and the connection to the norm bounds would change.

These details constrain transfer to ordinary training code. Updating both layers, introducing a bias, switching to multiclass cross entropy or using a deep convolutional backbone changes the problem. Those choices may be sensible in another experiment. They do not inherit this theorem automatically.

Our related discussion of pruning and generalization theory concerns a separate route to understanding capacity. This paper does not remove neurons from an existing model. It analyzes the width needed for a particular training guarantee.

Why nonsmooth ReLU needs a different argument

ReLU changes slope at zero. Away from that point, its behavior is linear on each side. At the transition, a second derivative based argument cannot simply proceed as though the activation were smooth everywhere. That creates an obstacle for some stability analyses that control the curvature of the network objective.

The source contrasts its approach with Taheri and Thrampoulidis in 2024, whose discussed result assumes smooth, Lipschitz activations and gives a risk bound in expectation. It also compares with Ji and Telgarsky in 2019, who studied ReLU classification using Rademacher complexity. Both comparisons are drawn from the paper’s own reference list.

The goal is not to declare one proof tool universally better. It is to obtain a sharper guarantee while retaining the actual ReLU activation. The new analysis uses the special structure of ReLU gates and the restricted region visited by gradient descent.

For context, our article on biased gradients and stability explores a different generalization argument. Here, the decisive step is a complexity estimate for a localized function class, rather than an analysis of how retraining changes after replacing a data point.

Gradient descent gives more than a decreasing objective

The algorithm uses the full training sample at every step. It follows the gradient of empirical risk, with an explicit choice of ReLU subgradient. At a zero preactivation, the indicator in Equation 3.2 of the source includes equality, so the chosen activation derivative is one.

Equation 3. The paper’s update rule. The learning rate is (eta), and the output coefficient (a_j) remains fixed.

The general loss assumptions require a convex, nonnegative, smooth loss as a function of signed output, together with a derivative controlled by the loss itself. Logistic loss satisfies these requirements. Smoothness of the loss does not make the network activation smooth, which is precisely why handling the gates matters.

Theorem 3 compares the training trajectory with a reference model. Its main quantity combines the reference model’s empirical risk over the training horizon and the squared distance from initialization. A useful reference model must therefore fit well without being arbitrarily far away.

Equation 4. The reference quantity and average empirical risk consequence of Theorem 3, under its loss, initialization, width and step conditions.

The same theorem controls the distance of the iterates from the reference point and the movement of individual neurons from initialization. The latter is measured by the maximum hidden weight displacement, rather than only the total distance across the entire network.

Equation 5. The two neighborhood controls used in the subsequent complexity analysis.

This gives a more specific story than saying that a network has many parameters. Gradient descent may explore a constrained part of parameter space. Generalization analysis can exploit that restriction rather than charging the learner for every function the architecture could represent.

The gates split the complexity into two manageable pieces

Imagine the initial preactivation of a neuron on a training example. If it is far from zero and the neuron moves only a little, its activation sign remains unchanged. Over that neighborhood, the change in output is linear in the weight displacement.

If the initial preactivation is close to zero, a small update may change the sign. Those neuron and example combinations are harder to analyze. The authors do not pretend they are linear. Instead, they bound their contribution separately and use the random initialization to control how numerous they can be.

Lemma 2 quantifies the number of initial preactivations inside a strip near zero. Lemma 9 then uses that control alongside the two neighborhood constraints. One contribution handles the relatively small set of unstable activation cases. The other retains almost linear behavior and uses the Euclidean reference constraint.

Rademacher complexity measures how well a function class can correlate with random signs placed on sample points. A class that fits arbitrary signs too easily can have a large statistical price. Restricting the class to the local region reached by training gives the analysis more information than a bound based on per neuron displacement alone.

The reference constraint is especially valuable because it keeps relationships between hidden weights in the estimate. Treating every neuron as an independent worst case contributor can discard that structure and produce a looser bound.

The proof also distinguishes the empirical reference quantity from a population counterpart used to define an appropriate function class. This is a technical safeguard. Choosing a favorable class after seeing the data and treating it as independent of that data would skip a necessary part of the statistical argument.

“We develop improved Rademacher complexity estimates for a hypothesis space motivated by our optimization analysis”Lei and colleagues, Section 6

Low loss changes the generalization calculation

An ordinary complexity bound can give a sample dependence resembling the inverse square root of sample size. An optimistic bound uses additional information about a smooth nonnegative loss, including the fact that low empirical loss can reduce the generalization penalty.

The paper combines its localized complexity estimate with such a result from Srebro and colleagues in 2010. That supporting work is cited in the source. The important sequence is optimization localization, improved complexity estimation and then a loss sensitive generalization argument.

Remark 19 makes a useful distinction. Applying the optimistic bound alone to the older complexity estimate can improve the sample dependence while leaving a worse margin dependence. The new localized estimate improves the margin factor as well. The final result therefore needs both ingredients, not merely the observation that training loss is small.

This explains why an article that only repeats the abstract misses the main lesson. The rate changes because the proof preserves two types of geometric control and treats activation changes selectively. More iterations by themselves are not the entire reason the bound improves.

The margin assumption supplies the favorable regime

NTK separability concerns the tangent features at initialization. These are the network output derivatives with respect to hidden weights. The assumption requires a unit parameter direction whose inner product with every label signed tangent feature is at least a positive margin.

Equation 6. Assumption 4, with the separator called (V) here to distinguish it from the reference model in the optimization bound.

The margin describes how strongly the tangent features separate the labels. The parameter (eta) controls how concentrated the separating direction can be in an individual neuron. A positive margin and a sufficiently distributed direction provide the starting geometry for the realization argument.

Under the additional width and parameter conditions in Lemma 16, moving from initialization in that direction constructs a reference network with arbitrarily small logistic loss. The required reference distance grows logarithmically with the reciprocal target loss, scaled by the reciprocal margin.

The existence of this direction should not be inferred from a low training loss at the end of an ordinary experiment. It is an assumption about the tangent features and specified norm constraints. A numerical candidate with positive minimum signed tangent score can be an informative diagnostic on a finite sample. It is not a population separability certificate or proof that all theorem conditions hold.

What the main risk result actually guarantees

Theorem 18 specializes the analysis to logistic loss and NTK separability. It balances the optimization horizon against the sample size by setting the learning rate times the number of updates equal to the number of training examples. Several width, margin and concentration conditions remain part of the statement.

Equation 7. Theorem 18, conditional on all its hypotheses. The tilde suppresses logarithmic factors.

The object bounded is the average population risk along the trajectory. It is not a standalone guarantee for the final iterate. It also is not the same object as the risk of a network obtained by averaging all its weights, since the nonlinear network makes parameter averaging a separate operation.

There are useful consequences, but they need the correct interpretation. A uniformly sampled iterate has expected risk equal to the average over the recorded trajectory, conditional on that trajectory. A predictor formed by averaging the network scores can use convexity of logistic loss to relate its risk to the average constituent risk. Neither observation silently changes the theorem into a last checkpoint statement.

The logarithms also matter. Saying that the rate is almost optimal acknowledges those factors and the dependence on confidence and other parameters. The notation is suitable for comparing asymptotic behavior. It does not provide an exact numerical error bar for a chosen dataset.

This is a comparison of bounds under the discussed assumptions, not a leaderboard. The newer theorem improves the sample size dependence relative to the listed earlier ReLU result while retaining a confidence statement. The smooth activation comparison already has the favorable risk order, but covers a different activation setting and reports an expectation guarantee.

Minimal width is not a universal neuron count

The title could tempt a reader to look for the smallest usable hidden layer. The paper does not provide a universal number or an empirical architecture search rule. Its result concerns sufficient width for the analysis under specified geometry and training conditions.

For the commonly discussed separator concentration of one over the square root of width, Remark 19 states a sufficient width dependence that suppresses logarithms but scales with the inverse eighth power of the margin. Equation 4.25 includes explicit logarithmic training horizon terms and substantial constants.

Equation 8. The margin dependence must accompany the shorthand description of polylogarithmic width.

A moderate sample dependence and a severe margin dependence can coexist. A proof that reduces the need for width growing polynomially with sample size may still impose a very large sufficient width when the margin is small. The asymptotic improvement is meaningful without becoming an immediately economical design prescription.

Nor is a sufficient condition a necessary one. If a smaller network works well, that does not contradict the theorem. It can mean the proof’s sufficient conditions are conservative, or the example has helpful structure not captured by the bound. Establishing the truly smallest possible width would require a different necessity argument.

How to explore the result without inventing experiments

The source has no benchmark experiment section, no reported dataset accuracy and no runtime comparison. Its evidence is the mathematical analysis and proofs. A responsible practical companion can implement the update rule and record quantities the argument makes interesting, while clearly separating that exercise from the paper’s findings.

The listing below does exactly that. It fixes the output coefficients, uses paired Gaussian initialization, normalizes artificial inputs and trains with full sample gradient descent. It records average empirical loss, evaluation loss on an independent artificial sample, neuron displacement and changes in activation patterns.

The average loss is recorded before each update to match the indexing of the theorem’s iterates. The final model evaluation is kept as a separate report. That small implementation decision prevents an off by one comparison from obscuring which mathematical quantity is being illustrated.

Activation changes can help visualize the proof intuition, but a measured fraction of changes is not the Rademacher complexity bound. Similarly, a recorded displacement is not a check of the reference radius unless the required reference model and its loss are available.

Our discussion of batch size and feature learning is a reminder that optimizer choices can change the object being studied. Replacing full sample updates here with minibatch SGD, Adam or momentum would require a fresh argument.

Limitations that determine where the result transfers

The architecture is shallow and fully connected. Extending the method to deep ReLU networks is an open direction named by the authors. Convolutional and recurrent structures are also proposed for future study rather than analyzed in the current result.

The data assumptions matter just as much. Independent samples, normalized inputs, binary labels and a suitable separability regime are part of the setting. Label noise or a distribution that does not admit the required tangent margin can invalidate the headline specialization even though a network remains trainable.

The confidence result is conditional on the stated hypotheses. It is more informative about tail behavior than an expectation bound at a comparable level, but the two forms are not interchangeable without considering constants, confidence dependence and assumptions. A high probability label alone does not make every comparison favorable.

The rates describe population loss, which cannot be measured exactly with a finite test split. Held out empirical loss is an estimate. It is useful for a demonstration but should not be relabelled as the theorem’s population quantity.

Finally, this paper does not offer a new deployment system or a replacement optimizer. Its contribution is an explanation of why the familiar update can succeed in a specified regime. The most immediate practical benefit is sharper reasoning about width, movement, activation changes and the scope of generalization claims.

A narrower function class gives a stronger explanation

The central achievement is a risk analysis that retains nonsmooth ReLU and improves the sample dependence in a favorable margin regime. The proof does not need to pretend the network explores its entire expressive capacity. It uses what gradient descent actually constrains.

The conceptual shift is to keep two neighborhoods in view. The maximum movement of a hidden neuron controls gate changes, while distance from a useful reference model supports a sharper complexity estimate. Each constraint does a different job.

The idea may inform analyses beyond this model, especially where an algorithm restricts the part of parameter space it visits. That is a research possibility, not an established extension to every modern architecture. The source’s own open questions keep the boundary clear.

The remaining limitations are substantial. Small margins can demand large sufficient widths, the favorable theorem is an average trajectory statement and general data may fail separability. These limits do not cancel the result. They identify the price of the guarantee.

Future progress could tighten margin dependence, address more architectures and connect comparable theory to wider data regimes. For now, the paper gives researchers a concrete lesson in reading neural network bounds. Preserve the activation assumptions, the trained parameters and the exact predictor whose risk is controlled.

The useful question is not simply how many parameters the network has. It is which functions the training path can reach, and what the data geometry allows those functions to learn.

PyTorch implementation of the analyzed update

This is an editorial implementation of the paper’s model and full sample update with logistic loss. It includes a custom ReLU derivative, fixed output weights, paired initialization, explicit gradient computation, training diagnostics, evaluation and an artificial smoke test. It is not authors’ experimental software, and its finite sample margin diagnostic does not certify the theorem.

The code passed a Python syntax check. An independent numerical calculation checked the explicit gradient against finite differences away from activation boundaries. PyTorch was unavailable here, so the included PyTorch smoke test was not executed. Run the listing in an environment with PyTorch before using its outputs.

"""Educational implementation of Eq. (3.2), Lei et al., JMLR 27(34), 2026.
Requires PyTorch. No experimental results or theorem verification are implied.
Only hidden weights train. Fixed paired output weights scale as 1/sqrt(m).
The zero-preactivation subgradient is one, as explicitly selected in the paper.
"""
import math
from dataclasses import dataclass
import torch
from torch import nn
from torch.nn import functional as F


class PaperReLU(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        ctx.save_for_backward(x)
        return x.clamp_min(0)

    @staticmethod
    def backward(ctx, grad_output):
        (x,) = ctx.saved_tensors
        return grad_output * (x >= 0).to(x.dtype)


class ShallowReLU(nn.Module):
    def __init__(self, input_dim, width, seed=42):
        super().__init__()
        if input_dim < 1 or width < 2 or width % 2:
            raise ValueError('input_dim positive, width positive and even required')
        generator = torch.Generator().manual_seed(seed)
        half = torch.randn(width // 2, input_dim, generator=generator,
                           dtype=torch.float64)
        signs = torch.randint(0, 2, (width // 2,), generator=generator)
        signs = (2 * signs - 1).to(torch.float64)
        self.W = nn.Parameter(torch.cat([half, half.clone()], dim=0))
        self.register_buffer('a', torch.cat([signs, -signs]) / math.sqrt(width))
        self.register_buffer('W_initial', self.W.detach().clone())

    def forward(self, x):
        return PaperReLU.apply(x @ self.W.T) @ self.a

    @torch.no_grad()
    def tangent_features(self, x):
        gate = (x @ self.W_initial.T >= 0).to(x.dtype)
        return (gate[:, :, None] * self.a[None, :, None] *
                x[:, None, :]).flatten(1)


def validate_data(x, y):
    if x.ndim != 2 or y.shape != (x.shape[0],) or x.shape[0] == 0:
        raise ValueError('Expected nonempty x[n,d] and y[n]')
    if not torch.isfinite(x).all() or not torch.isfinite(y).all():
        raise ValueError('Data must be finite')
    if not torch.all((y == -1) | (y == 1)):
        raise ValueError('Labels must be -1 or +1')
    if not torch.allclose(x.norm(dim=1), torch.ones_like(y), atol=1e-9, rtol=0):
        raise ValueError('Inputs must have unit Euclidean norm')


def logistic_loss(logits, y):
    return F.softplus(-y * logits).mean()


@torch.no_grad()
def explicit_gradient(model, x, y):
    """Eq. (3.2) uses ell'(u)=-sigmoid(-u), with all n examples."""
    logits = model(x)
    coefficient = -y * torch.sigmoid(-y * logits)
    gate = (x @ model.W.T >= 0).to(x.dtype)
    return ((coefficient[:, None] * gate * model.a[None, :]).T @ x) / len(x)


@torch.no_grad()
def evaluate(model, x, y):
    validate_data(x, y)
    score = model(x)
    prediction = torch.where(score >= 0, torch.ones_like(y), -torch.ones_like(y))
    return {'logistic_risk': logistic_loss(score, y).item(),
            'classification_error': (prediction != y).double().mean().item()}


@dataclass
class Run:
    empirical_risks: list
    test_risks: list
    frobenius_displacements: list
    max_neuron_displacements: list
    activation_change_fractions: list


def train_full_batch(model, x, y, steps, eta, test=None):
    validate_data(x, y)
    if steps < 1 or eta <= 0 or eta > 16 / 5:
        raise ValueError('Positive steps and 0 < eta <= 16/5 required')
    # This checks a necessary logistic step condition, NOT all theorem conditions.
    if test is not None:
        validate_data(*test)
    initial_gates = (x @ model.W_initial.T >= 0)
    run = Run([], [], [], [], [])
    for _ in range(steps):
        # Record W_t before updating: the theorem averages t=1,...,T.
        with torch.no_grad():
            run.empirical_risks.append(logistic_loss(model(x), y).item())
            if test is not None:
                run.test_risks.append(evaluate(model, *test)['logistic_risk'])
            shift = model.W - model.W_initial
            run.frobenius_displacements.append(shift.norm().item())
            run.max_neuron_displacements.append(shift.norm(dim=1).max().item())
            gates = (x @ model.W.T >= 0)
            run.activation_change_fractions.append(
                (gates != initial_gates).double().mean().item())
            model.W.sub_(eta * explicit_gradient(model, x, y))
    return run


@torch.no_grad()
def candidate_tangent_margin(model, x, y):
    """Finite-sample diagnostic, not an optimal or population margin certificate."""
    feature = model.tangent_features(x)
    signed = y[:, None] * feature
    direction = signed.mean(dim=0)
    size = direction.norm()
    if size <= 1e-12:
        return {'gamma_candidate': float('nan'), 'beta_candidate': float('nan')}
    direction = direction / size
    return {'gamma_candidate': (signed @ direction).min().item(),
            'beta_candidate': direction.reshape_as(model.W).norm(dim=1).max().item()}


def artificial_data(n, d, seed):
    """Unit-sphere points labelled by first-coordinate sign, for a demonstration."""
    generator = torch.Generator().manual_seed(seed)
    x = torch.randn(n, d, generator=generator, dtype=torch.float64)
    x = x / x.norm(dim=1, keepdim=True)
    y = torch.where(x[:, 0] >= 0, 1., -1.)
    return x, y


def smoke_test():
    x, y = artificial_data(48, 4, 7)
    test = artificial_data(128, 4, 8)
    model = ShallowReLU(4, 64, seed=9)
    assert torch.allclose(model(x), torch.zeros_like(y), atol=1e-12)
    assert list(dict(model.named_parameters())) == ['W']
    loss = logistic_loss(model(x), y)
    loss.backward()
    assert torch.allclose(model.W.grad, explicit_gradient(model, x, y), atol=1e-12)
    model.W.grad = None
    # The custom derivative at exactly zero agrees with the paper's >= gate.
    z = torch.zeros(1, dtype=torch.float64, requires_grad=True)
    PaperReLU.apply(z).sum().backward()
    assert z.grad.item() == 1.0
    a_initial = model.a.clone()
    W_initial = model.W.clone().detach()
    run = train_full_batch(model, x, y, steps=100, eta=0.5, test=test)
    assert torch.equal(model.a, a_initial)
    assert not torch.equal(model.W, W_initial)
    assert all(math.isfinite(v) for v in run.empirical_risks)
    assert run.empirical_risks[-1] < run.empirical_risks[0]
    print('Artificial demonstration, not a theorem certificate')
    print('Average iterate training risk:', sum(run.empirical_risks) / len(run.empirical_risks))
    print('Average iterate test risk:', sum(run.test_risks) / len(run.test_risks))
    print('Final model evaluation:', evaluate(model, *test))
    print('Candidate tangent margin:', candidate_tangent_margin(model, x, y))
    print('Last recorded max neuron displacement:', run.max_neuron_displacements[-1])


if __name__ == '__main__':
    smoke_test()

Frequently asked questions

What does minimal width mean in this paper?

It concerns sufficient width for a theoretical guarantee under stated conditions. It is not a universal neuron count or proof of the smallest network that can work.

Does the paper guarantee the final trained checkpoint?

The main NTK separability theorem bounds average population logistic risk along the trajectory. It does not directly state the same guarantee for the final iterate.

What is the NTK margin assumption?

A unit direction in hidden parameter space must separate every label signed tangent feature at initialization by a positive margin, with an additional bound on how concentrated that direction is in a neuron.

Does polylogarithmic width mean the network is always small?

No. The favorable sample dependence coexists with a potentially substantial inverse margin dependence. The sufficient width conditions and suppressed logarithms must be considered together.

Does this result apply to Adam or deep convolutional networks?

The stated analysis concerns full sample gradient descent on a shallow fully connected ReLU network with fixed output weights. Other optimizers and architectures require separate analysis.

Are the code outputs reported results from the paper?

No. The paper provides theory rather than benchmark experiments. The listing is an editorial artificial demonstration. Its syntax and an independent numerical gradient calculation were checked, but the PyTorch smoke test was not run here.

Related reading on aitrendblend.com