The Local Learning Coefficient: A Singularity-Aware Complexity Measure

At a local minimum of the population loss landscape there is a natural notion of complexity, given by the number of bits required to specify the minimum to within a tolerance $\epsilon$. This idea is well-known in the literature on minimum description length (Grünwald and Roos,, 2019) and was used by (Hochreiter and Schmidhuber,, 1997) in an attempt to quantify the complexity of a trained neural network. However, a correct treatment has to take into consideration the degeneracy of the geometry of the population loss $L$, as we now explain.

Consider a local minimum $w^{*}$ of the population loss $L$ and a closed ball $B(w^{*})$ centered on $w^{*}$ such that for all $w\in B(w^{*})$ we have $L(w)\geq L(w^{*})$. Given a tolerance $\epsilon>0$ we can consider the set of parameters $B(w^{*},\epsilon)=\{w\in B(w^{*})\mid L(w)-L(w^{*})<\epsilon\}$ whose loss is within the tolerance, the volume of which we define to be

$$V(\epsilon):=\operatorname{Vol}(B(w^{*},\epsilon))=\int_{B(w^{*},\epsilon)}dw\,.$$

(2)

The minimal number of bits to specify this set within the ball is

$$-\log_{2}(V(\epsilon)/\operatorname{Vol}(B(w^{*})))\,.$$

(3)

This can be taken as a measure of the complexity of the set of low loss parameters near $w^{*}$. However, as it stands this notion depends on $\epsilon$ and has no intrinsic meaning. Classically, this is addressed as follows: if a model is regular and thus $L(w)$ is locally quadratic around $w^{*}$, the volume satisfies a law of the form $V(\epsilon)\approx c\epsilon^{d/2},$ where $c$ is a constant that depends on the curvature of the basin around the local minimum $w^{*}$, and $d$ is the dimension of $W$.

This explains why $\frac{d}{2}$ is a valid measure of complexity in the regular case, albeit one that cannot distinguish $w^{*}$ from any other local minima. The curvature $c$ is less significant than the scaling exponent, but can distinguish local minima and $\log c$ is sometimes used as a complexity measure.

The population loss of a neural network is not locally quadratic near its local minima, from which it follows that the functional form of the volume $V(\epsilon)$ is more complicated than in the regular case. The correct functional form was discovered by Watanabe, (2009) and we adapt it here to a local neighborhood of a parameter (see Appendix A for details):

In the case where $m({w^{*}})=1$, the formula simplifies, and

$$V(\epsilon)\propto\epsilon^{{\lambda({w^{*}})}}$$

(5)

Thus, the LLC ${\lambda({w^{*}})}$ is the (asymptotic) volume scaling exponent near a minimum ${w^{*}}$ in the loss landscape: increasing the error tolerance by a factor of $a$ increases the volume by a factor of $a^{{\lambda({w^{*}})}}$. Applying Equation (3), the number of bits needed to specify $V(\epsilon)$ within $B(w^{*})$, for sufficiently small $\epsilon$ and in the case $m(w^{*})=1$, is approximated by

$$-{\lambda({w^{*}})}\log_{2}\epsilon+O(\log_{2}\log_{2}\epsilon)\,.$$

Informally, the LLC tells us the number of additional bits needed to halve an already small error of $\epsilon$:

$$-\log_{2}\Big{[}V(\tfrac{\epsilon}{2})/V(\epsilon)\Big{]}\approx-{\lambda({w^{ *}})}\log_{2}\frac{\epsilon}{2}+{\lambda({w^{*}})}\log_{2}\epsilon={\lambda({w ^{*}})}$$

See Figure LABEL:fig:volume_scaling for simple examples of the LLC as a scaling exponent. We note that, in contrast to the regular case, the scaling exponent ${\lambda({w^{*}})}$ depends on the underlying data distribution $q(x,y)$.

The (global) learning coefficient $\lambda$ was defined by Watanabe, (2001). If $w^{*}$ is taken to be a global minimum of the population loss $L(w)$, and the ball $B(w^{*})$ is taken to be the entire parameter space $W$, then we obtain the global learning coefficient $\lambda$ as the scaling exponent of the volume $V(\epsilon)$. The learning coefficient and related quantities like the WBIC (Watanabe,, 2013) have historically seen significant application in Bayesian model selection (e.g. Endo et al.,, 2020; Fontanesi et al.,, 2019; Hooten and Hobbs,, 2015; Sharma,, 2017; Kafashan et al.,, 2021; Semenova et al.,, 2020).

In Appendix C, we prove that the LLC is invariant to local diffeomorphism: that is, roughly, a locally smooth and invertible change of variables. This property is motivated by the desire that a good complexity measure should not be confounded by superficial differences in how a model is represented: two models which are essentially the same should have the same complexity.

Consider the following integral

$$Z_{n}(B_{\gamma}(w^{*}))=\int_{B_{\gamma}(w^{*})}\exp\{-nL_{n}(w)\}\varphi(w) \,dw,$$

(6)

where $L_{n}(w)$ is again the sample negative log likelihood, $B_{\gamma}(w^{*})$ denotes a small ball of radius $\gamma$ around ${w^{*}}$ and $\varphi$ is a prior over model parameters $w$. If (6) is high, there is high posterior concentration around ${w^{*}}$. In this sense, (6) is a measure of the concentration of low-loss solutions near ${w^{*}}$. Next consider a log transformation of $Z_{n}(B_{\gamma}(w^{*}))$ with a negative sign, i.e.,

$$F_{n}(B_{\gamma}(w^{*}))=-\log Z_{n}(B_{\gamma}(w^{*})).$$

(7)

This quantity is sometimes called (negative) log marginal likelihood or free energy, depending on the discipline. Given a local minimum ${w^{*}}$ of the population negative log likelihood $L(w)$, it can be shown using the machinery of SLT that, asymptotically in $n$, we have

$$F_{n}(B_{\gamma}(w^{*}))=n\underbrace{L_{n}({w^{*}})}_{\text{energy}}+ \underbrace{\lambda({w^{*}})}_{\text{entropy}}\log n+o_{p}(\log\log n),$$

(8)

where ${\lambda({w^{*}})}$ is the theoretical LLC expounded in Section 3. Remarkably, the asymptotic approximation in (8) holds even for singular models such as neural networks; further discussion on this can be found in Appendix B.

The asymptotic approximation in (8) suggests that a reasonable estimator of ${\lambda({w^{*}})}$ might come from re-arranging (8) to give what we call the idealized LLC estimator,

$$\hat{\lambda}^{\mathrm{idealized}}({w^{*}})=\frac{F_{n}(B_{\gamma}(w^{*}))-nL_ {n}({w^{*}})}{\log n}.$$

(9)

But as indicated by the name, the idealized LLC estimator cannot be easily implemented; computing or even MCMC sampling from the posterior to estimate $F_{n}(B_{\gamma}(w^{*}))$ is made no less challenging by the need to confine sampling to the neighborhood $B_{\gamma}(w^{*})$. In what follows, we use the idealized LLC estimator as inspiration for a practically-minded LLC estimator.

The first step towards a practically-minded LLC estimator is to circumvent the constraint posed by the neighborhood $B_{\gamma}(w^{*})$. To this end, we introduce a localizing Gaussian prior that acts as a surrogate for enforcing the domain of integration given by $B_{\gamma}({w^{*}})$. Specifically, let

$$\varphi_{\gamma}(w)\propto\exp\left\{-\frac{\gamma}{2}||w||^{2}_{2}\right\}$$

be a Gaussian prior centered at the origin with scale parameter $\gamma>0$. We replace (6) with

$$Z_{n}(w^{*},\gamma)=\int\exp\{-nL_{n}(w)\}\varphi_{\gamma}(w-w^{*})\,dw,$$

which, for $\beta=1$, can also be recognized as the normalizing constant to the posterior distribution given by

$$p(w|w^{*},\beta,\gamma)\propto\exp\left\{-n\beta L_{n}(w)-\frac{\gamma}{2}||w- w^{*}||_{2}^{2}\right\},$$

(10)

where $\beta>0$ plays the role of an inverse temperature. Large values of $\gamma$ force the posterior distribution in (10) to stay close to ${w^{*}}$. A word on the notation $p(w|w^{*},\beta,\gamma)$: this is a distribution in $w$ solely, the parameters ${w^{*}},\beta,\gamma$ are fixed, hence the normalizing constant to (10) is an integral over $w$ only. As $Z_{n}(w^{*},\gamma)$ is to be viewed as a proxy to (6), we shall accordingly treat

$$F_{n}(w^{*},\gamma)\coloneqq-\log Z_{n}(w^{*},\gamma)$$

as a proxy for $F_{n}(B_{\gamma}(w^{*}))$ in (7). Although it is tempting at this stage to simply drop $F_{n}(w^{*},\gamma)$ into the idealized LLC estimator in place of $F_{n}(B_{\gamma}({w^{*}}))$, we have to address estimation of $F_{n}(w^{*},\gamma)$, the subject of the next section.

Let us denote the expectation of a function $f(w)$ with respect to the posterior distribution in (10) as

$$\mathrm{E}_{w|{w^{*}},\beta,\gamma}f(w)\coloneqq\int f(w)p(w|w^{*},\beta, \gamma)\,dw.$$

Consider the quantity

$$\mathrm{E}_{w|{w^{*}},{\beta^{*}},\gamma}[nL_{n}(w)]$$

(11)

where the inverse temperature is deliberately set to $\beta^{*}=1/\log n$. The quantity in (11) may be regarded as a localized version of the widely applicable Bayesian information criterion (WBIC) first introduced in (Watanabe,, 2013). It can be shown that (11) is a good estimator of $F_{n}(w^{*},\gamma)$ in the following sense: the leading order terms of (11) match those of $F_{n}(w^{*},\gamma)$ when we perform an asymptotic expansion in sample size $n$. This justifies using (11) to estimate $F_{n}(w^{*},\gamma)$. Further discussion on this can be found in in Appendix D

Going back to (9), we approximate $F_{n}(B_{\gamma}(w^{*}))$ first with $F_{n}(w^{*},\gamma)$, which is further estimated by (11). We are finally ready to define the LLC estimator.

Note that $\hat{\lambda}(w^{*})$ depends on $\gamma$ but we have suppressed this in the notation. Let us ponder the pleasingly simple form that is the LLC estimator. The expectation term in (12) is a measure of the loss $L_{n}$ under perturbation near ${w^{*}}$. If the perturbed loss, under this expectation, is very close to $L_{n}({w^{*}})$, then ${\hat{\lambda}({w^{*}})}$ is small. This accords with our intuition that if ${w^{*}}$ is simple, its loss should not change too much under reasonable perturbations. Finally we note that in applications, we use the empirical loss $L_{n}(w)$ to determine a critical point of interest, i.e., ${\hat{w}_{n}^{*}}\coloneqq\operatorname*{arg\,min}_{w}L_{n}(w)\,.$ We lose something by plugging in ${\hat{w}_{n}^{*}}$ to (12) directly since we end up using the dataset $\mathcal{D}_{n}$ twice. However we do not observe adverse effects in our experiments, see Figure 4 for an example.

The LLC estimator defined in (12) is not prescriptive as to how the expectation with respect to the posterior distribution should actually be approximated. A wide array of MCMC techniques are possible. However, to be able to estimate the LLC at the scale of modern deep learning, we must look at efficiency. In practice, the computational bottleneck to implementing (12) is the MCMC sampler. In particular, traditional MCMC samplers must compute log-likelihood gradients across the entire training dataset, which is prohibitively expensive at modern dataset sizes.

If one modifies these samplers to take minibatch gradients instead of full-batch gradients, this results in stochastic-gradient MCMC, the prototypical example of which is (Welling and Teh,, 2011)’s Stochastic Gradient Langevin Dynamics (SGLD). The computational cost of this sampler is much lower: roughly the cost of a single SGD step times the number of samples required. The standard SGLD update applied to sampling (10) at the optimal temperature $\beta^{*}$ required for the LLC estimator is given by

$\displaystyle\Delta w_{t}$

$\displaystyle=\frac{\epsilon}{2}\left(\frac{\beta^{*}n}{m}\sum_{(x,y)\in B_{t} }\nabla\log p(y|x,w_{t})+\gamma(w^{*}-w_{t})\right)+N(0,\epsilon)$

where $B_{t}=\{(x_{i},y_{i})\}_{i=1}^{m}$ is a randomly sampled minibatch of samples of size $m$ for step $t$ and $\epsilon$ controls both step size and variance of injected Gaussian noise. Crucially, the log-likelihood gradient is evaluated using mini-batches. In practice, we choose to shuffle the dataset once and partition it into a sequence of size $m$ segments as minibatches instead of drawing fresh random samples of size $m$.

Let us now suppose we have obtained $T$ approximate samples $\{w_{1},w_{2},\dots,w_{T}\}$ of the tempered posterior distribution at inverse temperature $\beta^{*}$ via SGLD. This is usually taken from the SGLD trajectory after burn-in. We can then form what we call the SGLD-based LLC estimator,

$$\hat{\lambda}^{\mathrm{SGLD}}({w^{*}}):=n{\beta^{*}}\left[\frac{1}{T}\sum_{t=1 }^{T}L_{n}(w_{t})-L_{n}({w^{*}})\right].$$

(13)

For further computation saving, we also recycle the forward passes that compute $L_{m}(w_{t})$, which is required for computing $\nabla_{w}L_{m}(w_{t})$ via back-propagation, as unbiased estimate of $L_{n}(w_{t})$. Here by $L_{m}(w_{t})$, we mean $-\frac{1}{m}\sum_{(x,y)\in B_{t}}\log p(y|x,w_{t})$, though the notation suppresses the dependence on $B_{t}$ for brevity. Pseudocode for this minibatch version of the SGLD-based LLC estimator is provided in Appendix G. Henceforth, when we say the SGLD-based LLC estimator, we are referring to the minibatch version. In Appendix H, we give a comprehensive guide on best practices for implementing the SGLD-based LLC estimator including choices for $\gamma,\epsilon$, number of SGLD iterations, and required burn-in.

In this section, we verify the accuracy and scalability of our LLC estimator against theoretical LLC values in deep linear networks (DLNs) up to 100M parameters. Recall DLNs are fully-connected feedforward neural networks with identity activation function. The input-output behavior of a DLN is obviously trivial; it is equivalent to a single-layer linear network obtained by multiplying together the weight matrices. However, the geometry of such a model is highly non-trivial — in particular, the optimization dynamics and inductive biases of such networks have seen significant research interest (Saxe et al.,, 2013; Ji and Telgarsky,, 2018; Arora et al.,, 2018).

Thus one reason we chose to study DLN model complexity is because they have long served as an important sandbox for deep learning theory. Another key factor is the recent clarification of theoretical LLCs in DLNs by Aoyagi, (2024), making DLNs the most realistic setting where theoretical learning coefficients are available. The significance of Aoyagi, (2024) lies in the substantial technical difficulty of deriving theoretical global learning coefficients (and, by extension, theoretical LLCs) which means that these coefficients are generally unavailable except in a few exceptional cases, with most of the research conducted decades ago (Yamazaki and Watanabe, 2005b, ; Aoyagi et al.,, 2005; Yamazaki and Watanabe,, 2003; Yamazaki and Watanabe, 2005a, ). Further details and discussion of the theoretical result in Aoyagi, (2024) can be found in Appendix I.

The results are summarized in Figure 4. Our LLC estimator is able to accurately estimate the learning coefficient in DLNs up to 100M parameters. We further show that accuracy is maintained even if (as is typical in practice) one does not evaluate the LLC at a local minimum of the population loss, but is instead forced to use SGD to first find a local minimum $\hat{w}_{n}^{*}$. Further experimental details and results for this section are detailed in Appendix J. Overall it is quite remarkable that our LLC estimator, after the series of engineering steps described in Section 4, maintains such high levels of accuracy.

Here, we empirically test whether implicit regularization effectively induces a preference for simplicity as manifested by a lower LLC. We examine the effects of SGD learning rate, batch size, and momentum on the LLC. Note that we do so in isolation, i.e., we do not look at interactions between these factors. For instance, when we vary the learning rate, we hold everything else constant including batch size and momentum values.

We perform experiments on ResNet18 trained on CIFAR10 and show the results in Figure 1. We see that the training loss reaches zero in most instances and therefore on its own cannot distinguish between the effect of the implicit regularization. In contrast, we see that there is a consistent pattern of “stronger implicit regularization = higher test accuracy = lower LLC". Specifically, higher learning rate, lower batch size, higher momentum all apply stronger implicit regularization which is reflected in lower LLC. Full experimental details can be found in Appendix K. Note that in Figures 1 we employed SGD without momentum in the top two rows. We repeat the experiments in these top two rows for SGD with momentum; the associated results in Figure K.1 support very similar conclusions. We also conducted some explicit regularization experiments involving L2 regularization (Figure K.2 in Appendix K) and again conclude that stronger regularization is accompanied by lower LLC.

In contrast to the LLC estimates for DLN, the LLC estimates for ResNet18 cannot be calibrated as we do not know the true LLC values. With many models in practice, we find value in the LLC estimates by comparing their relative values between models with shared context. For instance, when comparing LLC values we might hold everything constant while varying one factor such as SGD batch size, learning rate, or momentum as we have done in this section.