An Analytical Theory of Auxiliary Learning

Federico Milanesio; Alessandro Ingrosso; Matteo Osella

Source-linked AI summary

An Analytical Theory of Auxiliary Learning

Federico Milanesio, Alessandro Ingrosso, Matteo Osella

arXiv:2609.29774v1Sep 24, 2026cs.LG

TL;DR

Although multi-task learning is used to improve performance, how auxiliary tasks affect generalization remains incompletely understood. This paper analyzes auxiliary learning in a teacher-student framework and finds relationships between main-task, auxiliary-task, and single-task generalization errors.

Problem
Despite the use of additional tasks to improve neural-network performance, how multi-task learning supports better generalization remains poorly understood.
Source: §1 · p. 1
Method
The authors analyze online training in a teacher-student framework, deriving dynamics and a leading-order linear error expression plus a nonlinear Lyapunov-based theory.
Source: §1 · pp. 1–2, §4 · p. 9
Results
For linear models, the optimal auxiliary-task weight is 2^-1 regardless of task correlation; for nonlinear activations, main and auxiliary errors follow a relationship involving the corresponding single-output error.
Source: §1 · p. 2, §4 · p. 9
Takeaways & Limitations
Within the studied setting, the theory predicts generalization error and the auxiliary-task weighting that minimizes it across activation functions and hyperparameters.
Source: §1 · p. 2, §4 · p. 9
The analysis leaves unexplored student models with a different number of neurons from the teacher.
Source: §4 · p. 9

Abstract

from arXiv · show

Auxiliary learning is an optimization paradigm in which a neural network's performance on a target task is improved by jointly training it on additional tasks. However, the mechanisms behind this improvement remain poorly understood. We study this problem using a teacher-student framework and derive a closed system of differential equations describing the dynamics of online stochastic gradient descent in the large-input limit. For linear networks, we obtain a closed-form expression for the generalization error to leading order in the learning rate, quantifying how task correlations and label noise determine the benefit of auxiliary learning. For non-linear activation functions, we develop a fluctuation-dissipation analytical theory that establishes a general relation linking the main and auxiliary errors to the corresponding single-task error. Numerical experiments support the theoretical predictions and show how auxiliary tasks improve generalization by balancing the forcing dynamics towards the optimal solution with gradient noise.

1 Introduction

The paper addresses the limited understanding of how auxiliary tasks can improve generalization despite competing objectives. It develops theory for their effects through training dynamics, building on multi-task learning, teacher–student, adaptive-filter, and fluctuation-dissipation work.

1 Introduction

Neural networks can overfit despite their capacity to approximate functions precisely, and multi-task learning may act as an inductive bias, but this mechanism remains poorly understood.The paper situates its question against empirical use of auxiliary objectives and the intuition that competing objectives might worsen performance.Source: §1 · p. 1
The paper generalizes mean-field ODEs for large-input teacher–student learning to auxiliary tasks, then derives a leading-order generalization-error approximation for multi-layer linear models.The framework studies training dynamics while leaving the target-task optimum unchanged.Source: §1 · pp. 1–2
For nonlinear activations, a Lyapunov-equation analysis predicts a relation between main-task, auxiliary-task, and corresponding single-output errors across activation choices.The authors describe this as extending a relationship previously recovered only for specific activations.Source: §1 · p. 2

1.1 Previous literature

Prior work uses mean-field ODEs for large-input teacher–student SGD, studies online linear-regression SGD as adaptive filtering, and relates SGD fluctuations to Lyapunov equations.The cited foundations include Saad & Solla (1995a,b), Biehl & Schwarze (1995), Widrow & Walach (1983), Haykin (1986), Sayed (2008), and Mandt et al. (2017).Source: §1.1 · p. 2

2 Auxiliary learning in the teacher-student framework

The paper formulates auxiliary learning as a two-output teacher–student problem and asks how to weight the auxiliary loss to minimize target generalization error. In the large-input limit, the dynamics reduce to closed ODEs over weight statistics.

2.1 Problem setup

The student has main and auxiliary outputs trained against targets from a same-sized teacher, and the auxiliary loss does not change the shared-teacher minimizer.The setup uses Gaussian inputs, a shared hidden layer, and training-label noise.Source: §2.1 · pp. 2–3
The central optimization question is how to choose the auxiliary-loss weight β to minimize target-task generalization error, primarily with the main-loss weight fixed at α = 1.The paper also explores α = 1 − β in Appendix E, where online-learning dynamics yield different results.Source: §2.1 · p. 3

2.2 The large-input limit

With Gaussian inputs, student and teacher pre-activations are jointly Gaussian, and weight dependence enters through student–student, student–teacher, and teacher–teacher overlaps.These overlap statistics are the problem's order parameters and determine both generalization error and dynamics.Source: §2.2 · p. 4
In the N → ∞ limit, the order parameters and second-layer weights follow a closed, deterministic set of ordinary differential equations for online SGD.Generalization error is computed from Gaussian averages of activation functions that depend on the order parameters.Source: §2.2 · p. 4

3 Results

The linear model yields a leading-order generalization-error result and shows that auxiliary-task weighting interacts with task correlation and learning dynamics. For nonlinear activations, empirical error relationships motivate a Lyapunov-based theory, whose predictions agree with simulations and whose scope includes a leading-order approximation.

3.1 Linear model

For equal-norm readouts, the deep linear model’s converged generalization loss has a leading-order-in-µ expression referenced to the auxiliary-free error.Here ρ is the correlation between the output weights, and ε_a.f. is the generalization error without an auxiliary task.Source: §3.1 · p. 4
With α = 1, the optimal β is independent of ρ; at ρ = 0 the main-task error is insensitive to β, while coincident tasks trade off effective learning rate and noise.With α = 1 − β, the effective learning rate is independent of β and the results differ; the appendix also reports greater auxiliary weighting when its noise is lower and gains from multiple auxiliary tasks.Source: §3.1 · pp. 5–6

3.2 Nonlinear activation functions

For erf and ReLU activations, the behavior is more complex and depends on more than ρ; at α = β = 1, the main-task error equals ε0 and a weighted-error relation approximately holds.The relation’s right-hand side is independent of the readout weights and their correlation, although both task errors depend on them.Source: §3.2 · pp. 6–7
Empirically, β* < 1/2 when α = 1; the general theory finds this bound for erf but not ReLU, except under a strong commuting assumption.The commuting condition does not hold generally.Source: §3.2 · p. 7, §3.3 · p. 9

3.3 A general analytic theory

At first order in µ, the stationary SGD fluctuations around W* are described by a Lyapunov equation that can be reduced from KN × KN to K × K and solved numerically.The resulting theoretical errors match the simulations and ODEs in Figure 3.Source: §3.3 · p. 7
The nonlinear theory recovers the single-output error using only J2, simplifying prior calculations but remaining limited to leading order in µ.The paper also focuses on well-converging initializations; a complete study of auxiliary learning’s possible role in escaping bad minima is beyond its scope.Source: §3.2 · p. 6, §3.3 · p. 8

4 Discussion and Future Work

For the teacher-student setup in the large-input, small-learning-rate limit, auxiliary learning improves performance; the discussion connects linear and nonlinear error behavior and identifies extensions left open.

4 Discussion and Future Work

The linear-model results are established in the large-input and small-learning-rate limit, and the discussion reports that auxiliary learning improves performance in this setting.The optimal linear auxiliary weight is independent of task correlation.Source: §4 · p. 9
The linear model’s optimal auxiliary weight is β*=1/2 regardless of task correlation, and its main and auxiliary errors relate to the corresponding single-output model’s generalization error.The nonlinear model’s optimal weight need not match the linear one, although its advantage curves remain similar and the error relationship is empirically valid.Source: §4 · p. 9
For nonlinear activations, expanding online SGD around the teacher optimum yields a Lyapunov equation that numerically predicts generalization error for choices of second-layer weights, activation, and hyperparameters.The theory also predicts the first-order term in µ for a one-output teacher-student network with fixed second-layer weights and any nonlinear activation, a novel result in that literature.Source: §4 · p. 9
Unequal teacher and student neuron counts remain unexplored; extending the theory to learnable second-layer weights or non-Gaussian inputs is proposed as future work.The authors note that the learnable-weight extension is not trivial because the effective learning rate varies between student weights.Source: §4 · p. 9

Contributions

The section records author contributions, generative-AI use, and code availability for reproducing numerical simulations.

Contributions

FM developed the theory with AI assistance and wrote the numerical-simulation code; MO conceived the original idea, and all authors revised the manuscript.The authors specify that “AI” here means the researcher, not a language model.Source: Contributions · p. 9

AI use statement

LLMs helped connect the work to dissipation-fluctuation theory and adaptive filters, while applying generic mathematical claims to this work was human-made.The authors also report using generative AI to aid software coding and review the manuscript, and take responsibility for the final content.Source: AI use statement · p. 9

Reproducibility statement

Code is available to reproduce the numerical simulations.Source: Reproducibility statement · p. 9

A Derivation of the auxiliary dynamics differential equations

The appendix derives auxiliary-task dynamics as mean-field ODEs and identifies the Gaussian averages needed to compute generalization error.

The derivation of the order-parameter differential equations follows prior one-output teacher-student analyses, including Saad & Solla (1995a,b), Biehl & Schwarze (1995), and Goldt et al. (2019).Source: A Derivation of the auxiliary dynamics… · p. 12
With i.i.d. Gaussian inputs, student and teacher pre-activations are jointly Gaussian; Q and R evolve during training, while T remains fixed.Source: A Derivation of the auxiliary dynamics… · p. 12
For the two-output case, the dynamics are obtained by taking expectations of the update equations and rescaling time as τ=t/N.The second-layer weights use a rescaled learning rate µ/N.Source: A Derivation of the auxiliary dynamics… · pp. 12–13
The dynamics form a non-extensive system of mean-field ODEs, and computing generalization error requires evaluating I2, a two-dimensional Gaussian average of a nonlinear function.The required integrals may be computed analytically or numerically.Source: A Derivation of the auxiliary dynamics… · pp. 12–13

A.1 Error function

The appendix gives integral calculations for the error function and ReLU activations, using analytic expressions where available and numerical quadrature for an intractable ReLU integral.

A.1 Error function

For the error-function case, the appendix states that all integrals can be computed analytically.Source: §A.1 · p. 13

A.2 Rectified linear unit

For ReLU, some integrals have analytic forms, but I4 cannot be computed analytically and is instead approximated using conditional Gaussian decomposition with Gauss–Hermite quadrature.A Hermite-expansion approximation was found to have convergence issues for ill-conditioned matrices.Source: §A.2 · p. 14
The ReLU calculation evaluates the bivariate normal CDF Φ2 using Owen’s identity, a smooth one-dimensional integral over a compact interval that converges rapidly under low-order Gauss–Legendre quadrature.I4 can then be approximated by Gauss–Hermite quadrature; when C is singular, it reduces to integrals known analytically.Source: §A.2 · p. 15

B Generalization error of the two-output linear model

The linear two-output analysis derives main-task generalization error under auxiliary training and examines how task correlations and label noise affect the optimal auxiliary weight. It also relates perfectly correlated outputs to single-output training with effective learning rate and noise.

B Generalization error of the two-output linear model

The optimal β* can exceed 1 and does not depend on task correlation ρ.Source: B Generalization error of the… · p. 18
When auxiliary-label noise exceeds main-task noise, a nonzero optimal auxiliary weight improves main-task generalization, though the benefit is marginal when δ > 1.Choosing β above the optimum can degrade performance.Source: B Generalization error of the… · p. 18
With two auxiliary labels of equal weight and uncorrelated outputs, β = γ = 1/2 yields better generalization than the minimum in Equation 9.Source: B Generalization error of the… · pp. 18–19

C Effective dynamics with perfectly correlated output

When both output weight vectors coincide, training is equivalent to a one-output teacher-student network with an effective learning rate and noise.Source: C Effective dynamics with perfectly… · p. 19

D.1 Derivation of the Lyapunov equation

The analysis approximates online SGD fluctuations around the teacher optimum with a Lyapunov equation, then simplifies its structure to study generalization and the optimal auxiliary weight. The theory is leading-order in learning rate: it agrees with ODE integration for small μ, while higher-order effects alter results as μ approaches 1.

D.1 Derivation of the Lyapunov equation

Linearizing the sample gradient around W* gives a stationary Lyapunov equation for the covariance of SGD fluctuations, using the population Hessian and gradient covariance at the optimum.The derivation assumes covariance scales as O(μ) and retains a cross-term because the sample Hessian and optimum gradient are correlated.Source: §D.1 · pp. 19–20

D.2 Simplifying the Lyapunov equation

Kronecker structure reduces the covariance equation from KN×KN to K×K, although the unreduced C in Equation 75 is still KN×KN.Source: §D.2 · p. 20

D.3 A consideration on the structure of P

At the fixed point, J2 = m²11T + κI; for a linear model κ = 0, making P rank one and reducing the problem to the span of v and u.Source: §D.3 · p. 21

D.4 A discussion of the optimal beta

When Pmain and Paux commute, joint diagonalization decouples the Lyapunov equation; in the one-mode case β* < 1/2, while the linear noncommuting case is exactly solvable.The linear-model solution matches the results in Appendix B.Source: §D.4 · p. 21

D.5 Effects of higher-order terms

Theory and ODE integration agree for small μ, but higher-order effects modify results as μ approaches 1.Figure S2 compares the relative error between the Lyapunov-based theory and integrated ODEs across μ.Source: §D.5 · p. 21, Figure S2 · p. 22

D.6 Convergence of different initializations

The section compares how often ODE integrations from different initializations converge to the theoretical generalization error. The reported comparison finds more convergence with an auxiliary output than without one.

D.6 Convergence of different initializations

The preliminary exploration reports that not all initializations of W converge to a neighborhood of W* at the end of the dynamics.It samples W* from N(0, 1) and checks how often the ODE dynamics error matches theoretical generalization obtained by numerically solving the corresponding Lyapunov equation.Source: §D.6 · pp. 22–23
More initializations converge to the global minimum with β = 1/2 than with no auxiliary output (β = 0); McNemar’s exact test rejects equal marginal convergence probabilities (p < 10^-10).The authors suggest auxiliary learning can help escape bad minima in this teacher-student setup.Source: §D.6 · p. 23, Table S1 · p. 23
In the teacher-student setup, McNemar’s exact test rejected equal marginal convergence probabilities for the two β values (p < 10−10), suggesting auxiliary learning can help escape bad minima.The analysis could be expanded to more β values, though the authors do not expect major differences.Source: §D.6 · p. 23

E A different scaling of the main task

With α=1−β, changing the auxiliary weight also changes the main-task learning rate when tasks are not perfectly correlated, producing different online-learning behavior. For nonlinear activations, the reported observations place the optimal β between 1/2 and 1.

With α=1−β, increasing β reduces the main-task learning rate when tasks are not perfectly correlated, and β* is at least 1/2.For this weighting, β must remain below 1.Source: E A different scaling of the main task · p. 23
For ρ=1, β*=1/2: equal task weights still improve generalization because the noise on each task is independent.Source: E A different scaling of the main task · p. 23
For erf and ReLU activations, behavior depends on more than the correlation ρ between output weights.Figure S4 compares relative main-task error across ρ and β for linear, erf, and ReLU models.Source: E A different scaling of the main task · p. 23, Figure S4 · p. 24
For nonlinear models, α=β=1/2 yields εmain=εaux=ε0/2, while αεmain+βεaux=ε0(α²+β²).Source: E A different scaling of the main task · p. 23, Figure S4 · p. 24