PaperlayerView exact revision on arXiv

Source-linked AI summary

Path-conditioned training: a principled way to rescale ReLU neural networks

Arthur Lebeurrier, Titouan Vayer, Rémi Gribonval

arXiv:2602.19799v2stat.MLcs.LGmath.OC

TL;DR

ReLU rescalings preserve a network's function but can yield different training dynamics, and principled ways to exploit this symmetry remain limited. The paper uses path-lifting to choose a geometry-based rescaling that conditions a kernel in lifted space. On CIFAR-10, it reaches baseline accuracy in fewer epochs.

  • Problem

    Principled ways to leverage ReLU rescaling symmetries to improve training dynamics remain limited, even though rescaling-equivalent parameters can lead to different training trajectories.

    Source: §1 · pp. 1–2

  • Method

    PathCond uses path-lifting to select a rescaling that aligns the path kernel with a reference, aiming to improve conditioning in lifted space.

    Source: §1 · p. 2, §4.2 · p. 4

  • Results

    On CIFAR-10, PathCond reaches baseline accuracy in 1.5× fewer epochs using rescaling at initialization alone.

    Source: §1 · p. 2

  • Takeaways & Limitations

    PathCond offers a geometry-based approach to rescaling ReLU networks while preserving their implemented function.

    Source: §1 · p. 2

    The equivalence of the optimization to the standard log-determinant formulation is stated for positive-definite Gθ, corresponding to full-rank ∂Φ.

    Source: §4.3 · p. 5

Abstract

from arXiv · show

Despite recent algorithmic advances, we still lack principled ways to leverage the well-documented rescaling symmetries in ReLU neural network parameters. While two properly rescaled weights implement the same function, the training dynamics can be dramatically different. To offer a fresh perspective on exploiting this phenomenon, we build on the recent path-lifting framework, which provides a compact factorization of ReLU networks. We introduce a geometrically motivated criterion to rescale neural network parameters which minimization leads to a conditioning strategy that aligns a kernel in the path-lifting space with a chosen reference. We derive an efficient algorithm to perform this alignment. In the context of random network initialization, we analyze how the architecture and the initialization scale jointly impact the output of the proposed method. Numerical experiments illustrate its potential to speed up training.

1. Introduction

ReLU networks have function-preserving rescaling symmetries, but these symmetries can affect training dynamics. The paper uses path-lifting to motivate a principled rescaling method and reports faster training on CIFAR-10.

1. Introduction

  • The paper frames its motivation as improving understanding of optimization dynamics and translating those insights into practical improvements for state-of-the-art models.The introduction also notes the computational challenge of training large language models.Source: §1 · p. 1
  • ReLU rescaling preserves the implemented function, yet the paper identifies principled use of this symmetry to improve training dynamics as an open problem.Prior approaches are described as often guided more by practical considerations than formal theoretical principles.Source: §1 · p. 1
  • Path-lifting provides an intermediate representation between network parameters and the functions they implement, which the paper uses to develop a geometry-aware rescaling criterion.The criterion moves parameters to a rescaling-equivalent choice intended to improve training properties.Source: §1 · p. 2
  • On CIFAR-10, PathCond reaches baseline accuracy in 1.5× fewer epochs using rescaling at initialization alone.The paper describes the criterion as parameter-free and reports that it does not degrade training on the evaluated settings.Source: §1 · p. 2

2. Sketch of the idea on a simple example

A one-neuron ReLU example shows how rescaling-equivalent parameters can produce different training trajectories. PathCond chooses a rescaling intended to make lifted-space gradient descent resemble direct gradient descent in that space.

2. Sketch of the idea on a simple example

  • For the one-neuron model, rescaling incoming parameters and inversely rescaling the output parameter preserves the function, while the path-lifting representation is unchanged.The example uses fθ(x) = u ReLU(vx + w) and positive rescaling factors.Source: §2 · p. 2
  • In the stated data example, the loss factors through Φ(θ), and matching the target lifted representation implements ReLU with zero training loss.This motivates viewing learning in lifted space as well as parameter space.Source: §2 · p. 2
  • PathCond seeks a rescaling that aligns the path kernel with the identity, approximating lifted-space gradient flow without pseudo-inverting large matrices.The path kernel varies under rescaling even though Φ(θ) remains invariant.Source: §2 · p. 2
  • In the toy example, PathCond-rescaled initializations converge faster and their lifted-space trajectories more closely follow the idealized gradient-descent trajectory.The figure contrasts three initializations and shows loss, lifted-space trajectories, and parameter-space trajectories.Source: §2 · p. 3, Figure 1 · p. 3
  • The paper connects its approach to NTK conditioning, whose spectrum is described as governing training-loss convergence rates in particular regimes.For certain over-parameterized lazy regimes, larger NTK eigenvalues lead to faster convergence rates.Source: §2 · p. 3

3. Rescaling symmetries and the path-lifting

ReLU networks admit function-preserving rescalings, but gradient flow generally depends on which rescaling-equivalent parameters initialize training. Path-lifting represents products of weights along paths, factoring out this parameter redundancy and enabling a local loss factorization.

3. Rescaling symmetries and the path-lifting

  • Rescaling a neuron's incoming weights and bias by λ and outgoing weights by 1/λ preserves its input-output mapping and the network function.The paper defines parameter vectors related by admissible diagonal rescalings as rescaling equivalent.Source: §3 · p. 3
  • Gradient flow generally follows different trajectories from rescaling-equivalent initial parameters, and extreme imbalances can lead to arbitrarily poor optimization and generalization performance.The paper presents this sensitivity as an opportunity to guide training through controlled rescaling.Source: §3 · pp. 3–4
  • Path-lifting maps parameters to a representation whose path coordinates are products of weights along network paths, capturing essential geometry while eliminating redundant degrees of freedom.Rescaling-equivalent parameters share the same lifted representation.Source: §3 · p. 4
  • The lifted representation induces a local factorization of the global loss, enabling analysis of dynamics in lifted space rather than redundant parameter space.The passage attributes this factorization to the path-lifting framework.Source: §3 · p. 4

4. Path-conditioned training

The section treats parameter rescaling as a way to shape optimization geometry in path-lifting space, then develops PathCond to align the path kernel with identity using an efficient logdet-based procedure.

4. Path-conditioned training

  • In lifted space, the path kernel acts as a local metric and preconditions the flow induced by parameter-space gradient flow.Rescaling leaves the lifted representation unchanged but changes the path kernel, thereby shaping the optimization dynamics.Source: §4 · p. 4, §4.1 · p. 4

4.1. Rescaling is pre-conditioning

  • Rescaling is a preconditioning choice because it changes the path-kernel geometry while preserving the network’s lifted representation.The section notes that several intuitive parameter- or Hessian-based criteria agree on toy examples but not on larger models.Source: §4.1 · p. 4

4.2. Proposed rescaling criterion

  • PathCond chooses rescalings to align the path kernel with identity, avoiding explicit formation of the combinatorially large path-space kernel.The criterion targets the kernel rather than directly optimizing alignment between path-space gradients, whose dimension grows with the number of paths.Source: §4.2 · pp. 4–5

4.3. Explicit algorithm with the logdet divergence

  • A generalized-determinant variant handles zero eigenvalues; equivalence with standard logdet is stated when G is positive definite.The stated condition is that ∂Φ is full rank.Source: §4.3 · p. 5
  • The logdet criterion reduces the optimization to neuron-dimensional variables and requires only diag(Gθ), computable with one backward pass.The optimization dimension is the number of hidden neurons, typically smaller than the parameter count.Source: §4.3 · p. 5
  • Alternating minimization with closed-form coordinate updates yields the explicit algorithm, with time O(niter p) and space O(p + H).Typical convergence takes fewer than 10 iterations.Source: §4.3 · p. 6

5. Experiments

Across CIFAR-10 and MNIST experiments, PathCond rescaling at initialization accelerates training without degrading generalization. The experiments also examine its computational cost and identify architectural and initialization regimes associated with stronger effects.

5.1. Faster Training Dynamics

  • At small learning rates, PathCond matches or exceeds baseline performance, while its effect becomes less pronounced as the learning rate increases.At higher rates, results remain comparable to the baseline and ENorm.Source: §5.1 · p. 6
  • On CIFAR-10 MLPs, PathCond reaches target accuracy up to 1.5× faster; unlike ENorm, its criterion has no hyperparameter.The comparison spans network depths; ENorm does not yield improvements in this experiment.Source: §5.1 · p. 6, §5 · p. 6, Figure 2 · p. 7

5.2. Generalization

  • On CIFAR-NV, PathCond reaches 80% best test accuracy versus 77% for Baseline and ENorm, about 20 epochs earlier in near-optimal training performance.The figure plots training loss, training accuracy, and test accuracy.Source: §5.2 · pp. 6–7, Figure 3 · p. 7

5.3. Computational Cost

  • The rescaling criterion uses only diag(G), costs O(# params) to compute, and can be evaluated with a single backward pass.The block coordinate descent algorithm for the resulting optimization is also described as efficient.Source: §5.3 · p. 7

5.4. What are the Favorable Regimes for PathCond ?

  • On MNIST autoencoders, PathCond achieves smaller final training loss in every tested case, with a more pronounced effect at larger compression factors.The figure also reports the maximum absolute log rescaling at initialization across compression factors.Source: Figure 5 · p. 8, §5.4 · p. 9

5.3. Computational Cost

  • On an NVIDIA RTX A4000, rescaling costs 22%–103% of one CIFAR-10 training epoch across tested architectures.Because rescaling is performed once at initialization, its overhead over a typical 100+ epoch run is described as negligible.Source: Figure 6 · p. 8, §5.3 · p. 8

5.4. What are the Favorable Regimes for PathCond ?

  • PathCond is favored when diag(G) varies widely, as with irregular layer widths or small initialization variance in near-constant-width networks; constant diag(G) yields identity rescaling.Figure 4 validates that rescaling magnitude grows as widths become more irregular and initialization variance decreases.Source: §5.4 · pp. 8–9, Figure 4 · p. 8

6. Conclusion and discussions

The section summarizes PathCond’s geometric, path-lifting approach to rescaling ReLU networks for improved gradient-flow conditioning, alongside its optimization and implementation properties, including efficient computation of the diagonal of G and measured rescaling costs. It reports approximate input-output preservation and training acceleration with GELU and SiLU, and identifies broader optimizer and architecture extensions as future work.

6. Conclusion and discussions

  • PathCond aims to improve gradient-flow conditioning in path-lifting space; with GELU and SiLU, rescaling remains approximately input-output preserving and still accelerates training.The authors suggest combining PathCond with adaptive optimizers such as Adam, Muon, or Shampoo, and extending its analysis and procedure to gated activations and transformer architectures.Source: §6 · p. 9

A. Properties of the path-lifting

B. Gradient in Φ-space

  • The pseudo-inverse preconditioner gives maximum alignment with the lifted-space gradient, but computing it requires a large p × p matrix that PathCond circumvents.Source: B. Gradient in Φ-space · p. 13

C. Equivalence between different rescaling criteria on toy examples

  • In the two-weight toy example, when a ≥ 0, increasing the rescaled parameter norm increases the Hessian operator norm and condition number.The passage notes that a ≥ 0 holds, for example, with quadratic loss.Source: C. Equivalence between different… · p. 14

D. Characterization of the rescaling symmetry

E. Bregman matrix divergences

  • For spectral Bregman divergences, the nonzero-eigenvalue divergence for AA⊤ depends only on its nonzero eigenvalues; the logdet case uses the generalized determinant.Source: E. Bregman matrix divergences · p. 19

F. PathCond optimization problem

  • The PathCond optimization problem uses the path kernel Pθ, the parameter-space Gram matrix Gθ, and the logdet spectral divergence in the q ≥ p regime.The passage says q ≥ p applies to most DAG ReLU networks, except very small ones.Source: F. PathCond optimization problem · p. 19
  • PathCond represents rescaling with neuron coordinates u via D′=diag(exp(Bu)); if every Gii>0, the objective has a solution and each coordinate update has a polynomial-root step.The coordinate step is given by the unique positive root when g>0.Source: F. PathCond optimization problem · pp. 20, 22, 23

G. Computation of the diagonal of G

  • The diagonal of G is computed efficiently using differentiation with respect to squared parameters θ²=θ⊙θ.Here G=∂Φ⊤∂Φ, and θ² denotes the element-wise square of the parameter vector.Source: G. Computation of the diagonal of G · p. 23

H. Detailed Algorithm and Complexity Analysis

  • PathCond stores only nonzero entries of B and incrementally updates Bu and a global auxiliary sum, giving O(niter·p) total complexity.Each coordinate update touches only the incoming and outgoing parameters of its neuron.Source: H. Detailed Algorithm and Complexity… · pp. 23–24

H.1. Computational Cost: Additional Results

  • 98% of PathCond rescaling time is spent in block coordinate descent, and total rescaling time scales linearly with parameter count.The reported timing covers MLPs with varying depth and constant hidden width.Source: §H.1 · p. 24, Figure 7 · p. 25

H.2. Number of BCD Iterations: Convergence Analysis

The BCD convergence rate varies by architecture: the MLP approaches its final objective decrease quickly, while FullyConv and ResNet-18C need more iterations. The authors use 10 iterations as a compromise between solution quality and computational cost.

Experimental setup

  • The ablation monitors normalized objective decrease and mean per-neuron primal updates, with small updates indicating BCD convergence.The normalized objective is in [0, 1] and equals 0 at convergence; the mean update is the average per-neuron change between iterations.Source: §H.2 · p. 25

Results

  • The MLP achieves about 90% of its final objective decrease within 5 iterations, while FullyConv and ResNet-18C need 25–30 iterations to approach convergence.At 10 iterations, the latter architectures have realized only 20–30% of their total objective decrease.Source: §H.2 · p. 25
  • Using niter = 10 is a practical compromise: it gives near-optimal solutions for simpler architectures and substantial improvements for more complex ones at small computational overhead.The reported overhead is typically half to one CIFAR-10 epoch across architectures.Source: §H.2 · p. 25

I. Architecture Statistics: Number of Parameters, Hidden Units, and Paths

This appendix defines the architecture statistics used for the experiments and analyzes expected path magnitudes under standard initialization. It also identifies the constant-diagonal case, where the rescaling algorithm has no effect.

I. Architecture Statistics: Number of Parameters, Hidden Units, and Paths

  • The architecture statistics include trainable parameters, hidden units, and input-to-output computational paths; the path count is computed by a forward pass with all parameters set to one.The path count follows from the path-1 norm when each path product equals one.Source: §I · pp. 25–26

J. Regimes Analysis at Initialization

J.1. The Matrix G

  • The analysis focuses on diag(G), with a parameter-specific diagonal entry defined in terms of paths ending at output neurons.Source: §J.1 · p. 27

J.2. Layered Fully-Connected ReLU Networks

  • For layered fully-connected ReLU networks, expected path magnitudes are analyzed by counting paths through each parameter and computing expected squared magnitudes for path types; final-layer biases are omitted.Source: §J.2 · pp. 27–28

J.3. Classical Initializations

  • Under standard initialization, Proposition J.2 gives a layer-dependent expected path magnitude for parameters in layer k.The stated initialization assumption includes He or Kaiming initialization with a constant variance-scaling factor.Source: §J.3 · p. 29

J.4. Consequences for the Rescaling Algorithm

The expected variation in path magnitudes, and thus the expected impact of rescaling, depends on both layer widths and initialization scale. The analysis identifies varying-width networks and near-constant-width networks with small initialization variance as especially favorable regimes.

J.4. Consequences for the Rescaling Algorithm

  • At large variance, expected path magnitudes depend primarily on nL/nk+1, with ratios across layers of order nk′+1/nk+1 when a is sufficiently large.Source: §J.4 · p. 29
  • At small variance, expected path magnitudes can differ across layers by a factor of order 1/a, producing strong variation.Source: §J.4 · p. 29
  • For nonconstant-width architectures, the expected diagonal of G varies across parameters, so rescaling can have a nontrivial effect.Source: §J.4 · p. 29

J.5. Constant-Width Networks

  • For constant-width networks at a = 1, expected path magnitude increases linearly with depth, but the diagonal is approximately constant and rescaling has little effect when n ≫ L.Source: §J.5 · p. 30
  • For constant-width networks with a ≠ 1, the diagonal has a multilevel structure with significant variation across layers, especially in deeper layers.Source: §J.5 · p. 30

J.6. Summary and Implications

  • The strongest expected effects are in varying-width networks at any initialization scale and near-constant-width networks initialized with small standard deviation relative to 1/√n.CIFAR-10 input and output dimensions create architectural irregularity despite intermediate layers of width 500; gains decrease as depth increases.Source: §J.6 · p. 31
  • In these favorable regimes, variation in expected path magnitudes produces ill-conditioning of G that the rescaling algorithm can address.Source: §J.6 · p. 31

K. Additional Experiments for Section 5.1

PathCond performs strongly at small to moderate learning rates, but performance degrades at larger rates. The reported explanation is that optimization instabilities then overshadow initialization benefits, outside the small-learning-rate assumptions.

Additional experiments

L. Additional Experiments for Section 5.2

Across learning rates from 10^-2 to 10^-4, PathCond improves convergence and generalization at small to moderate learning rates. Its advantage narrows as learning rates rise, and the smallest-rate run uses a longer training schedule.

L. Additional Experiments for Section 5.2

M. Additional Experiments for Section 5.4

Supplementary experiments examine how learning rate relates to architecture width variation and rescaling adjustments. Required rescaling magnitude increases with compression factor across all three learning rates.

M. Additional Experiments for Section 5.4

N. ENorm : Experimental Setup

The ENorm comparisons use a fixed default configuration. For smooth activations, PathCond no longer preserves the network mapping exactly, but initialization-time rescaling still yields a qualitative training speed-up in the tested MLPs.

N. ENorm : Experimental Setup

  • ENorm uses p = 2, c = 1, one rescaling cycle per SGD iteration, and excludes BatchNorm parameters.This is the stated default configuration; p = 2 minimizes summed squared layer ℓ2 norms, while c = 1 removes the depth penalty.Source: N. ENorm : Experimental Setup · pp. 32–33

O. Smooth Activation Functions

  • PathCond exactly preserves the input-output mapping for positively 1-homogeneous activations such as ReLU, but not for smooth GELU or SiLU.For GELU and SiLU, the rescaling remains approximately input-output preserving at initialization.Source: O. Smooth Activation Functions · p. 33
  • GELU and SiLU approximate ReLU on positive inputs, while SiLU has a slightly larger negative dip.Source: O. Smooth Activation Functions · p. 33
  • The activation study replaces ReLU with GELU or SiLU in MLPs with hidden width 500 and depths 2 through 9, then measures relative output difference after initialization rescaling.Source: O. Smooth Activation Functions · p. 33
  • The approximation-error estimate averages over 10 initializations and input samples; constraining the primal iterate to ∥u∥∞≤5 bounds rescaling factors to [e^-5, e^5].Without this constraint, the rescaling can produce NaNs in GELU MLPs.Source: O. Smooth Activation Functions · p. 33
  • For GELU, relative output error decreases from approximately 10^-1 at L = 2 to 10^-3 at L = 9, while PathCond still yields the qualitative training speed-up seen with ReLU.Figure 13 compares PathCond performance across network depths for GELU MLPs.Source: O. Smooth Activation Functions · p. 34, Figure 13 · p. 34
Loading 2602.19799v2…