Source-linked AI summary
A Flatness-Generalization Relation in the Teacher-Student Tree-Committee Machine
Brandon Livio Annesi, Davide Straziota, Enrico Maria Malatesta
TL;DR
Evidence for a flatness–generalization relation is contested, and its validity has not been established from Hessian-spectrum analyses in this setting. The paper studies the relation in teacher–student tree committee machines by characterizing ERM minimizers and their Hessian spectra. It finds that correlations depend on the task, phase, and flatness measure, and can reverse in classification.
Problem
The flatness–generalization relation is a widely used but contested heuristic, and prior theoretical Hessian-spectrum studies had not examined it from this perspective.
Source: §I · pp. 1–2Method
The paper uses replica methods to characterize typical ERM minimizers and Edwards-Jones formalism to derive their limiting Hessian spectrum and flatness measures.
Source: §I · p. 2Results
Flatness–generalization correlations depend on task and phase: in regression the spectral mean and right edge correlate with test error, while classification correlations hold for all three measures only when overparametrized and can reverse in recovery.
Source: §VI · p. 9Takeaways & Limitations
Flatness is not a universal indicator of generalization; its relation to test error depends on the observable, learning task, and phase.
Source: §I · p. 2, §VI · p. 9The framework does not explain the origin of finite-number left outliers at large α or their relation to class-structured Hessian outliers.
Source: §VI · p. 9
Abstract
from arXiv · showhide
The flatness of the loss landscape at a minimizer is a widely used heuristic for reasoning about neural-network generalization, yet evidence for this relation is mostly empirical and controversial. We study this relation in a teacher-student tree committee machine, where both the ERM estimator and the Hessian spectrum are analytically tractable in the proportional high-dimensional limit. First, we use a zero-temperature Gibbs formulation to obtain predictions for the observables of the typical minimizers of the empirical loss. Secondly, we use Edwards-Jones formalism to derive the limiting Hessian resolvent around these typical minimizers. All predictions agree with finite-size gradient-descent simulations. Finally, we study three measures of flatness, namely the left and right edges and the spectral mean, and check if a decrease in generalization error as the dataset size is increased corresponds to an increase in flatness. We find that the answer strongly depends on the learning task and on the ratio of the number of parameters to the number of data points. In regression, the spectral mean and right edge correlate with the generalization error, while the left edge does so only in the overparametrized regime. In classification this correlation reliably holds only in the highly overparametrized phase, while for underparametrized networks it can even reverse.
I. INTRODUCTION
Flatness is a widely used but contested heuristic for understanding generalization. The paper studies the relation analytically in teacher–student tree committee machines and finds that it depends on the observable, task, and phase.
I. INTRODUCTION
- The flatness–generalization relation is contested: prior work reports counterexamples and reparameterizations that vary Hessian-based measures without changing network outputs.Other work has proposed rescaling-invariant flatness measures and compared the flatness of minimizers with different generalization properties.Source: §I · p. 1
- The Hessian’s largest eigenvalue measures sharpest-direction curvature, while its trace measures average curvature over parameter-space directions.Prior studies also associate Hessian spectral structure with training dynamics, batch size, robustness, and data structure.Source: §I · p. 1
- Committee machines offer an analytically tractable setting with non-convexity and hidden-unit specialization for studying learning phenomena.Teacher–student models permit macroscopic observables to be characterized using a small number of overlaps.Source: §I · p. 2
- The study characterizes typical empirical-risk minimizers and derives their Hessian spectrum using replica and Edwards-Jones formalisms.It finds RS or FRSB behavior depending on task, dataset size, dimensionality, and regularization, and uses 1RSB as a numerically tractable approximation to FRSB observables.Source: §I · p. 2
- The flatness–generalization relation varies by measure, task, and phase: regression and classification show different correlation patterns.For regression, the right edge and spectral mean correlate with test error, while classification correlations hold across all three measures only in the overparametrized phase.Source: §I · p. 2
- The paper reports good agreement between its theoretical predictions and gradient-descent simulations, including in the replica-symmetry-broken regime.The passage notes that exact equilibrium sampling in that regime is challenging for stable algorithms in related mean-field spin-glass models.Source: §I · p. 2
II. THE MODEL
The model trains a tree committee student on labels generated by a teacher of the same architecture. The student learns first-layer weights by minimizing a regularized empirical loss, and the zero-temperature Gibbs measure concentrates on ERM estimators.
II. THE MODEL
- The teacher and student are tree committee machines with K hidden units, non-overlapping receptive fields, and trainable first-layer weights.The dataset contains P = αN input–label pairs; teacher weights are unknown, and the teacher and student share second-layer weights.Source: §II · pp. 2–3
- The teacher generates labels using the identity output for regression and the sign output for binary classification.The student has the same architecture and learns its first-layer weights.Source: §II · pp. 2–3
- The student learns by minimizing the ℓ2-regularized empirical loss, with the data-fitting and regularization terms both of order N as α varies.The ERM set is studied through a Gibbs measure that concentrates on ERM estimators as β →∞.Source: §II · p. 3
- Generalization is measured by test mean-squared error in regression and teacher–student disagreement probability in classification.Both quantities are evaluated on a fresh Gaussian input.Source: §II · p. 3
A. Replica description and generalization error
The replica description expresses typical ERM properties through a finite set of order parameters. These parameters characterize overlaps and norms and determine generalization errors in the infinite-width limit.
A. Replica description and generalization error
- The replica calculation expresses the typical free entropy in terms of order parameters describing Gibbs-sample overlaps and alignment with the teacher.The analysis takes N, P →∞ at fixed α = P/N, then sends K →∞.Source: A. Replica description and… · p. 4
- The distribution of student-pair overlaps may be nontrivial, while the norm and teacher–student overlap have delta-function distributions.More elaborate symmetry-breaking ansatzes add overlap scales, with full RSB giving a continuous overlap distribution.Source: A. Replica description and… · pp. 4–5
- Under replica symmetry, independently sampled students have a single typical mutual overlap; 1RSB instead assigns distinct within-cluster and between-cluster overlaps.The 1RSB distribution uses q1 > q0, with Parisi parameter m encoding the probability of the between-cluster overlap q0.Source: A. Replica description and… · p. 4
- In the infinite-width limit, regression and classification errors become explicit functions of the order parameters Q and r.The errors are computed from the joint Gaussian law of teacher and student outputs on a fresh test input.Source: A. Replica description and… · p. 5
- The analysis focuses on square loss because it makes a Gaussian integral analytically evaluable; generic losses require numerical evaluation of that integral.The model also uses φ(x) = erf(x), allowing the NNGP kernel to be evaluated analytically.Source: A. Replica description and… · p. 5
B. Regression: a replica symmetric phase diagram
In regression, the full phase diagram is captured by a replica-symmetric solution. At low enough regularization, test error approaches zero for α > 1, and gradient-descent simulations closely match the theory.
B. Regression: a replica symmetric phase diagram
- The regression phase diagram is replica symmetric throughout, despite neither convexity nor Bayesian optimality guaranteeing replica symmetry here.The teacher output function is the identity, so teacher and student architectures are fully matched.Source: B. Regression: a replica symmetric… · p. 5
- For α > 1, sufficiently low regularization brings test error close to zero and nearly recovers the teacher; for α < 1, test error varies non-monotonically with λ.For α > 1, test error decreases monotonically as λ decreases.Source: B. Regression: a replica symmetric… · p. 6
- Gradient-descent experiments show excellent agreement with the replica-symmetric theory for regression.Figure 2 plots test error, squared norm Q, and teacher–student overlap r against sample complexity α for different regularization strengths λ; its markers are finite-size simulations.Source: B. Regression: a replica symmetric… · p. 6
C. Classification: Replica Symmetry Breaking
Classification moves from an RS regime at low sample complexity into an intermediate RSB phase, then returns to a qualitatively different RS recovery phase above a regularization-dependent threshold.
C. Classification: Replica Symmetry Breaking
- At low sample complexity, classification is described by an RS saddle point for any regularization strength λ.Source: C. Classification: Replica Symmetry… · p. 6
- At sufficiently small λ, increasing α opens an intermediate RSB window, approximated by 1RSB with only small corrections from higher RSB levels.The exact description requires full replica symmetry breaking.Source: C. Classification: Replica Symmetry… · p. 6
- Above a threshold depending on λ, the system exits RSB and returns to an RS recovery phase with a qualitatively different character from the low-α RS phase.Source: C. Classification: Replica Symmetry… · p. 6
- For sufficiently small regularization, the RSB-to-recovery transition is discontinuous with an order-parameter jump and coexistence, whereas larger regularization yields a continuous transition.The example λ = 0.001 has a discontinuous transition.Source: C. Classification: Replica Symmetry… · p. 6
IV. HESSIAN SPECTRUM IN MINIMIZERS
The limiting Hessian spectrum around typical minimizers is derived through its Stieltjes transform; for square loss, that transform obeys a cubic equation determined by the learning-state order parameters.
IV. HESSIAN SPECTRUM IN MINIMIZERS
- The calculation targets the limiting spectral distribution by computing the Stieltjes transform and recovering the distribution through Stieltjes inversion.A Gaussian integral representation of the determinant enables use of the replica method.Source: §IV · pp. 6–7
- For square loss, the Hessian Stieltjes transform satisfies a cubic equation whose coefficients depend on α, the activation kernel and its derivatives, and the learning order parameters.The spectrum is therefore determined by the same macroscopic state variables used to describe learning.Source: §IV · p. 7
- Theory and finite-N gradient-descent spectra agree excellently in RS phases, with a small spectral-shape discrepancy in RSB phases.At high α, simulations also show K teacher-aligned eigenvalues detaching from the left bulk edge; their vanishing spectral weight K/N excludes them from the limiting density.Source: §IV · p. 7
V. FLATNESS IS PHASE-DEPENDENT
The relationship between Hessian flatness and generalization varies by task and phase: regression shows selective correlations, while classification correlations hold reliably only in its first overparametrized RS phase.
V. FLATNESS IS PHASE-DEPENDENT
- The spectral mean measures average curvature, the right edge probes the sharpest direction, and the left edge probes soft or marginal directions.The analysis varies α and checks whether these observables correlate positively with generalization error.Source: §V · p. 8
- In regression, the spectral mean and right edge correlate positively with test loss across explored regularizations, while the left edge correlates only for α < 1 and becomes anticorrelated when underparametrized.The left-edge correlation strengthens with higher λ within the overparametrized regime.Source: §V · p. 8
- In classification, spectral observables correlate reliably with generalization only in the first overparametrized RS phase; they are essentially uncorrelated in RSB and can reverse sign in recovery.Figure 6 contrasts RS branches, 1RSB branches, and gradient-descent markers.Source: §V · p. 8, Figure 6 · p. 8
- Throughout the RSB phase, the left edge is fixed at −λ, saturating the lower bound imposed by positive semidefiniteness of the full-loss Hessian and indicating marginal states.The predicted marginality is matched by simulations.Source: §V · p. 8
VI. DISCUSSION AND LIMITATIONS
The study finds that flatness–generalization relations depend on task and phase, and its conclusions are validated against finite-size gradient-descent simulations.
VI. DISCUSSION AND LIMITATIONS
- The order parameters describing empirical-risk minimizers determine both their learning performance and the limiting Hessian spectrum, and predictions agree quantitatively with finite-size gradient-descent simulations.Source: §VI · pp. 8–9
- The central conclusion is that flatness–generalization relations are phase-dependent rather than universal, with regression and classification exhibiting different correlation patterns.In regression, the spectral mean and right edge track test error across explored regularizations, while the left edge does so only for α < 1; in classification, correlations hold only in the overparametrized phase and may reverse in recovery.Source: §VI · p. 9
- In the RSB phase, the left edge saturates the −λ lower bound and indicates marginal minimizers; at large α, finitely many left outliers are invisible to the limiting spectral density.The origin of these outliers and their relation to class-structured Hessian outliers in deep classifiers lie outside this framework.Source: §VI · p. 9
- The analysis varies the dataset-size-to-dimension ratio across landscapes, unlike the common fixed-dataset comparison across minimizers of one landscape.Whether the same phase-conditioned picture holds within a single landscape requires a fixed-α analysis, left to future work.Source: §VI · p. 9
VII. AI DISCLOSURE
The supplemental material records the authors’ use of generative AI and develops the replica-based analysis from the model setup through symmetry-breaking approximations and their low-temperature limit.
VII. AI DISCLOSURE
- Generative AI assisted small steps in the replica computation, some plots, and manuscript polishing; the authors say they independently verified and reviewed the work.They take responsibility for the final content, including AI-assisted text, claims, or artifacts.Source: §VII · p. 9
Appendix A: Setting and objective
- The setting uses identical one-hidden-layer tree committee machines, with regression or classification teachers, Gaussian inputs and teacher weights, and L2-regularized empirical risk minimization.The ERM estimator is studied through the large-β limit of the Gibbs measure.Source: Appendix A · p. 12
Appendix B: Equilibrium measure
- The replica calculation expresses equilibrium through overlap matrices whose entries encode student–student overlap, student squared norm, and teacher–student alignment.The teacher norm is one in the large-N limit.Source: Appendix B · p. 13
1. Large width limit
- The large-width limit is taken after P,N→∞, with K remaining smaller in order; independent input portions across hidden units simplify the overlap structure and free-entropy terms.The appendix identifies the resulting covariance with the NNGP kernel at initialization in the infinite-width limit.Source: §1 · p. 13
2. RS ansatz
- The RS ansatz imposes a symmetric structure on the overlap matrix to evaluate the equilibrium free entropy.The appendix then performs the n→0 limit.Source: §2 · p. 14
3. k-steps Replica Symmetry Breaking
- The k-step RSB ansatz replaces the RS overlap structure with nested blocks, and its k=0 and k=1 cases recover the RS and 1RSB approximations.The free entropy is optimized over overlap variables, breaking points, the squared norm, and teacher–student overlap.Source: §3 · pp. 14–15
a. Large β limit
- In the low-temperature limit, the order parameters and breaking points scale with β so the free energy remains finite, yielding corresponding limiting entropic and energetic terms.The appendix gives a modified recursion for the energetic term.Source: a. Large β limit · pp. 15–16
b. Square Loss case
This section develops a replica-based analysis of the square-loss case, covering solution norms, task-dependent generalization measures, and the Hessian spectrum. It presents the saddle-point and zero-temperature scaling equations alongside the resolvent equation and its spectral edges.
b. Square Loss case
- At large β, the squared norm Q changes substantially from the k = 0 RS solution to k = 1 RSB, while k = 20 adds comparatively small corrections.In the large-β comparison across α, more than one RSB step only slightly influences Q, and other order parameters remain essentially unchanged.Source: Figure 7 · p. 18, b. Square Loss case · pp. 18–19, Figure 9 · p. 19
- The overlap q versus breaking point m is shown alongside validation of the large-β scalings introduced in (B21).Figure 8 also compares overlap and breaking points after rescaling the breaking points as βm.Source: b. Square Loss case · p. 18, Figure 8 · p. 18
4. Generalization and training error
- Classification performance is measured by the probability of misclassifying a new teacher-generated Gaussian input, whereas regression performance is measured by test loss.The passages state that these quantities are computed in the K → ∞ limit using the replica method.Source: §4 · p. 19
1. Gibbs formulation of the spectrum of the Hessian
- The Hessian spectrum is computed for typical weights sampled from the Gibbs measure; quadratic L2 regularization shifts it, and the appendix spectrum requires rescaling by α.The stated relation is ρ_H(z) = α ρ_αH(αz).Source: §1 · pp. 20–21
2. Replica tricks; 3. Annealed ansatz
- The resolvent calculation uses replica tricks and an annealed ansatz for new overlaps, while the usual overlap parameters satisfy the equilibrium saddle-point equations.The Stieltjes transform is obtained from the saddle-point solution for P.Source: §2 · p. 21, §3 · pp. 23–24
5. Square loss case
- For square loss, Gaussian integrals yield the free entropy, and the saddle-point equation for P in both classification and regression is cubic.The coefficients differ between tasks because the constants C and T depend on whether the teacher task is classification or regression.Source: §5 · p. 26
- The cubic equation gives either three real roots with vanishing spectral density or a complex-conjugate pair whose negative-imaginary-part root is physical.Using this root yields the Hessian resolvent and the spectral density, and evaluating the equation at αz gives equation (16) of the main text.Source: §5 · p. 26
6. Zero temperature limit
- In the zero-temperature limit, effective order parameters follow β scalings, while the other effective order parameters remain finite.Taylor expansion expresses δΦ and δΣ through the generalized kernel function; the replica-symmetric case follows by setting ΦQ,q0 and ΣQ,q0 to zero.Source: §6 · pp. 26–27
a. Square loss case
- For square loss, the limiting calculation changes the coefficients A1 and A2, while task-dependent constants C and T have finite large-β limits.The constants are specified separately for classification and regression.Source: a. Square loss case · p. 27
7. Left and right edges of the spectrum of the Hessian
- The left and right spectral edges are obtained from the two solutions of dz/dP = 0, where z(P) ceases to be defined for real P beyond the extrema.Source: §7 · p. 28
8. A comment on the validity of the Annealed Ansatz
- The replica-symmetric saddle-point equations yield T = t = o = p = 0, confirming that the annealed ansatz is correct.The conjugate parameters are solved in terms of the order parameters; the free entropy is quadratic in T, t, and o, and p = 0 is always a solution.Source: §8 · p. 28
Appendix D: Details of the numerical experiments
The numerical experiments validate the theory with full-batch gradient descent on the Section II model. They evaluate the unregularized-loss Hessian by exact diagonalization of the full matrix.
Appendix D: Details of the numerical experiments
- Simulations use N = 2000, K = 20, and φ = erf, with 20 teachers per α and an independent test set of size 10000.For each α, the training set has P = ⌊αN⌋ i.i.d. standard Gaussian inputs.Source: Appendix D · p. 28
- Full-batch gradient-descent simulations (N = 2000, K = 20, φ = erf) minimize the regularized objective (3) with nominal weight decay λ and learning rate η0 ∈[0.5, 1.0], for up to Emax ∈[10⁴, 10⁵] epochs.Early stopping occurs when relative changes in the loss and order parameters Q = ∥w∥²/N and r = w · w⋆/N all fall below ε = 10⁻⁸.Source: Appendix D · p. 28
- At training end, the unregularized-loss Hessian is evaluated in closed form and the full N × N matrix is diagonalized exactly, without low-rank or stochastic approximation.The experiments use Julia, Flux.jl, ForwardDiff.jl, and LinearAlgebra.jl.Source: Appendix D · p. 28