A short 'derivation' of Watanabe's Free Energy Formula

post by Wuschel Schulz (wuschel-schulz) · 2024-01-29T23:41:44.203Z · LW · GW · 6 comments

Contents

  Background: Statistical Learning Theory
    The Probabilistic Model
    The True Distribution
    The Parameters
    The Data
    The Negative Log-Likelihood
    Prior over the Parameters
    Posterior Distribution
  Free Energy
None
6 comments

Epistemic status: I wrote the post mostly for myself, in the process of understanding the theory behind Singular Learning Theory, based on the SLT low lecture series. This is a proof sketch with the thoroughness level of an experimental physics lecture: enough to get the intuitions but consult someone else for details.

Background: Statistical Learning Theory

In the framework of statistical learning theory, we aim to establish a model that can predict an outcome $y$ given an input $x$ with a certain level of confidence. This prediction is typically governed by a set of parameters $w$ , which need to be learned from the data. Let us begin by defining the primary components of this framework.

The Probabilistic Model

The probabilistic model  $p(y|x,w)$ represents the likelihood of observing the outcome  $y$ given the input $x$ and the parameters $w$. This model is parametric, meaning that the form of the probability distribution is predefined, and the specific behavior of the model is determined by the parameters $w$.

The True Distribution

The true distribution $q\left(y|x\right)$  is the actual distribution from which the data points are sampled. In practice, this distribution is not known and  $p(y|x,w)$ attempts to approximate it through the learning process.

The Parameters

The parameters $w$  are the weights or coefficients within our model that determine its behavior. The goal of learning is to find the values of $w$ that make $p(y|x,w)$ as close as possible to the true distribution $q\left(y|x\right)$.

The Data

The data $D=\{x_i,y_i\}$  consist of pairs of input values $x_i$  and their corresponding outcomes $y_i$. These data points are used to train the model, adjusting the parameters $w$ to fit the observed outcomes.

The Negative Log-Likelihood

The negative log-likelihood is a measure of how well our model $p(y|x,w)$ fits the data $D$ . It is defined as the negative sum of the logarithms of the model probabilities assigned to the true outcomes $y_i$:

$NLL\left(w\right)=-{\sum }_{i=1}^N\log p(y_i|x_i,w),$

where $N$ is the number of data points in $D$. Minimizing the negative log-likelihood is equivalent to maximizing the likelihood of the data under the model, which is a central objective in the training of probabilistic models.

Prior over the Parameters

In a Bayesian framework, we incorporate our prior beliefs about the parameters $w$  before observing the data through a prior distribution $\phi \left(w\right)$. This prior distribution encapsulates our assumptions about the values that $w$ can take based on prior knowledge or intuition. For instance, a common choice is a Gaussian distribution, which encodes a preference for smaller (in magnitude) parameter values, promoting smoother model functions.

Posterior Distribution

Once we have observed the data  $D$, we can update our belief about the parameters $w$ using Bayes' theorem. The updated belief is captured by the posterior distribution, which combines the likelihood of the data given the parameters with our prior beliefs about the parameters. The posterior distribution is defined as:

$P\left(w|D\right)=\frac{1}{Z_n}\exp (-NLL(w\left)\right)\Phi \left(w\right),$

where  $NLL\left(w\right)$ is the negative log-likelihood of the parameters given the data, $\Phi \left(w\right)$ is the prior over the parameters, and $Z_n$ is the normalization constant, also known as the partition function. The partition function ensures that the posterior distribution is a valid probability distribution by integrating (or summing) over all possible values of $w$:

$Z_n=\int \exp (-NLL(w\left)\right)\phi \left(w\right)dw.$

The posterior distribution reflects how likely different parameter values are after taking into account the observed data and our prior beliefs.

Free Energy

This formulation looks like the Boltzmann distribution from statistical physics. This distribution also describes a probability: the probability that a system can be found in a particular microstate is $e^{-E}$ with $E$ being the energy (here assuming, that the temperature is 1). This formalism can be expanded to macrostates, when changing the Energy to Free Energy, taking into account that microstates that have many realizations are more likely.
We can do a similar step with the Posterior of our model weights, by considering neighborhoods around local minima as macrostates. This opens up the question of what the correct formula for the Free Energy of these macrostates are, so we can make predictions about the behavior of learning machines.
First, let us consider the case of large datasets. In that case, we can replace the negative Log-Likelihood with n times its average:

$\begin{array}{cc}NLL\left(w\right)& =-\sum _{i=1}^N\log p(y_i|x_i,w)\\ & \approx n\int \log \left(p\right(y|x,w\left)\right)q\left(y|x\right)q\left(x\right)dx\\ & =n(\int \log (\frac{p(y|x,w)}{q\left(y|x\right)}\left)q\right(y|x\left)q\right(x)dx-\int \log (q\left(y|x\right)\left)q\right(y|x\left)q\right(x\left)dx\right)\\ & =n<KL\left(p_w|q\right)-S\left(q\right){>}_{q\left(x\right)}\end{array}$

We see that we can compose the negative Log likelihood into the number of data points n, the expected Kullback–Leibler divergence of the model with the true distribution, and the expected entropy of the true distribution. As we see here, the expected entropy term does not contribute to the posterior of the parameters:
$\begin{array}{cc}P\left(w|n\right)& =\frac{1}{Z_n}\exp (-NLL(w\left)\right)\Phi \left(w\right)\\ & =\frac{\exp \left(\right(-NLL\left(w\right)\left)\Phi \right(w)}{\int \exp (-NLL(w\left)\right)\Phi \left(w\right)dw}\\ & =\frac{\exp (-n⟨KL(p_w\parallel q)-S(q\left){⟩}_{q\left(x\right)}\right)\Phi \left(w\right)}{\int \exp (-n⟨KL(p_w\parallel q)-S(q\left){⟩}_{q\left(x\right)}\right)\Phi \left(w\right)dw}\\ & =\frac{\exp (-n⟨KL(p_w\parallel q\left){⟩}_{q\left(x\right)}\right)\Phi \left(w\right)}{\int \exp (-n⟨KL(p_w\parallel q\left){⟩}_{q\left(x\right)}\right)\Phi \left(w\right)dw}\end{array}$
The macrostates, aka the regions in parameter space we are now interested in analyzing are the vicinities of local minima in the Loss landscape. Let us assume we decompose our parameter space into macrostates $W={\cup }_{\alpha }{W}_{\alpha }$.
$\begin{array}{cc}P(w\in W_{\alpha })& ={\int }_{w\in W_{\alpha }}p\left(w|n\right)dw\\ & =\frac{{\int }_{w\in W_{\alpha }}\frac{1}{Z_n}\exp (-L(w\left)\right)\Phi \left(w\right)dw}{{\sum }_{\alpha }{\int }_{w\in W_{\alpha }}\frac{1}{Z_n}\exp (-L(w\left)\right)\Phi \left(w\right)dw}\\ & =\frac{{\int }_{w\in W_{\alpha }}\exp (-n⟨KL(p_w\parallel q\left){⟩}_{q\left(x\right)}\right)\Phi \left(w\right)dw}{{\sum }_{\alpha }{\int }_{w\in W_{\alpha }}\exp (-n⟨KL(p_w\parallel q\left){⟩}_{q\left(x\right)}\right)\Phi \left(w\right)dw}\\ & =\frac{\exp -F(W_{\alpha },n)}{{\sum }_{\alpha }\exp -F(W_{\alpha },n)}\\ F(W_{\alpha },n)& =-\log \left({\int }_{w\in W_{\alpha }}\exp \right(-n⟨KL(p_w\parallel q){⟩}_{q\left(x\right)}\left)\Phi \right(w\left)dw\right)\end{array}$

We call this $F$ the free energy of this area, as it serves the same function as the Free Energy in statistical physics. We have constructed $W_{\alpha }$ such that it only contains one minimum of the expected Kulberg Leibler divergence (from here on  we abbreviate $KL(p_w\parallel q){⟩}_{q\left(x\right)}$ as $K\left(\omega \right)$). These minima don't have to be points but can be different sets  within $W_{\alpha }$.

 Since we assume, that $K\left(\omega \right)$ is smooth, it can locally be approximated by its Taylor series at any point. How the leading order terms of the Taylor series look at its minimum depends on the geometry of the minimum. Let's look at an example of how that might look like from the figure above. In a) we have a simple minimum point, so the Taylor series of $K$ at this point might look like $x^2+y^2$, going up in each direction. In b) our minimum is a one-dimensional submanifold, and so the Taylor expansion of $K$ will look like $x^0+y^2$, increasing in one dimension, but staying flat when going along another. The situation in c) looks almost the same, except in the center, where the line crosses itself. There the Taylor expansion will look like $x^2\cdot y^2$, locally staying flat in either direction. Points like these, where the directions you can go in while staying on the set suddenly change, are called singularities.
Let us now calculate the Free Energy around a point ${\omega }^{\ast }$, where $K$ has a minimum. In general the leading order term of the $K$ at a minimum looks like this:
$\begin{array}{c}K({\omega }^{\ast }+x)=K_{min}+a\prod _i{x}_i^{2k_i}:=K_{min}+a{\overrightarrow{x}}^{2\overrightarrow{k}}\end{array}$
Here $K_{min}$ is the value of the KL-divergence at its minimum.
To get it into this form, we have to make a change of coordinates for $x$ (for example when we are at a singularity where the different parts of the minimum come in from an angle). A process called blowup guarantees us, that we are always able to find such coordinates. Since we are at a minimum, we know that all leading order exponents are even.
Let's also approximate our prior over the weights  $\Phi (w^{\ast }+x)$ on this point in the same basis as its leading order term:
$\begin{array}{c}\Phi \left(\omega \right)=b\prod _i{x}_i^{h_i}:=b{\overrightarrow{x}}^{\overrightarrow{h}}\end{array}$
For large $n$, the Free Energy at any point far away from the minima gets exponentially suppressed. To see how much the minimal point ${\omega }^{\ast }$ contributes to the integral we can integrate it in its vicinity. Here we choose to integrate from 0 to 1. 

Why not from $-1$ to $1$?^[1]
$\begin{array}{c}F({\omega }^{\ast },n)=-log\left({\int }_0^1\exp \right(-n(K_{min}+a{\overrightarrow{x}}^{2\overrightarrow{k}})\left(b{\overrightarrow{x}}^{\overrightarrow{h}}d\overrightarrow{x}\right)\end{array}$
The first thing we can do is pull out the $K_{min}$ term as it does not depend on $x$.
$\begin{array}{c}F({\omega }^{\ast },n)=n{K}_{min}-log\left({\int }_0^1\exp \right(-na{\overrightarrow{x}}^{2\overrightarrow{k}}\left)b{\overrightarrow{x}}^{\overrightarrow{h}}d\overrightarrow{x}\right)\end{array}$
To solve the remaining integral, we first do a Laplace and then a Mellin transform, as this brings the integral into a form, that is easier to solve.
$\begin{array}{cc}L\left({\int }_0^1\exp \right(-na{\overrightarrow{x}}^{2\overrightarrow{k}}\left)b{\overrightarrow{x}}^{\overrightarrow{h}}d\overrightarrow{x}\right)& ={\int }_0^{\infty }\left({\int }_0^1\exp \right(-na{\overrightarrow{x}}^{2\overrightarrow{k}}\left)b{\overrightarrow{x}}^{\overrightarrow{h}}d\overrightarrow{x}\right)e^{-nt}dt\\ & ={\int }_0^1\delta (t-a{\overrightarrow{x}}^{2\overrightarrow{k}})b{\overrightarrow{x}}^{\overrightarrow{h}}d\overrightarrow{x}\\ M\left(L\left({\int }_0^1\exp \right(-na{\overrightarrow{x}}^{2\overrightarrow{k}}\left)b{\overrightarrow{x}}^{\overrightarrow{h}}d\overrightarrow{x}\right)\right)& ={\int }_0^{\infty }\left({\int }_0^1\delta \right(t-a{\overrightarrow{x}}^{2\overrightarrow{k}}\left)b{\overrightarrow{x}}^{\overrightarrow{h}}d\overrightarrow{x}\right)t^{z}dt\\ & ={\int }_0^1(a{\overrightarrow{x}}^{2\overrightarrow{k}}{)}^{z}b{\overrightarrow{x}}^{\overrightarrow{h}}d\overrightarrow{x}\\ & =a^{z}b\prod _i{\int }_0^1{x}_i^{2z{k}_i+h_i}\\ & =a^{z}b\frac{1}{{\prod }_i(2z{k}_i+h_i+1)}\end{array}$

Next, we group the dimensions $i$ along, such that within each group $I_{\alpha }$, the value ${\lambda }_{\alpha }=\frac{h_i+1}{2k_i}$ is the same. Later in the derivation, it will turn out, that only the biggest $\lambda$ will contribute for large $n$. 
$\begin{array}{cc}M\left(L\left({\int }_0^1\exp \right(-na{\overrightarrow{x}}^{2\overrightarrow{k}}\left)b{\overrightarrow{x}}^{\overrightarrow{h}}d\overrightarrow{x}\right)\right)& =a^{z}b\frac{1}{{\prod }_{\alpha }{\prod }_{i\in I_{\alpha }}2k_i(z+{\lambda }_i)}\end{array}$
Now we go backward, and reverse the Mellin and the Laplace transform:
$\begin{array}{cc}L\left({\int }_0^1\exp \right(-na{\overrightarrow{x}}^{2\overrightarrow{k}}\left)b{\overrightarrow{x}}^{\overrightarrow{h}}d\overrightarrow{x}\right)& ={\int }_{-i\infty }^{+i\infty }\frac{a^{z}b}{2\pi i}\frac{1}{{\prod }_{\alpha }{\prod }_{i\in I_{\alpha }}2k_i(z+{\lambda }_i)}{t}^{-z-1}dz\\ & =\frac{b}{2\pi it}{\int }_{-i\infty }^{+i\infty }\frac{e^{z\log \left(\frac{a}{t}\right)}}{{\prod }_{\alpha }{\prod }_{i\in I_{\alpha }}2k_i(z+{\lambda }_i)}dz\end{array}$
We can solve this integral with the Residue Theorem, closing the integral path over the negative half of the if $a>t$. Since all ${\lambda }_{\alpha }$ have to be positive, when we close the integral over the negative half, we enclose no residues and the integral comes out as zero. Otherwise each unique ${\lambda }_i$, which appears $m$ times, being a residual of degree $m$:
$\begin{array}{cc}L\left({\int }_0^1\exp \right(-na{\overrightarrow{x}}^{2\overrightarrow{k}}\left)b{\overrightarrow{x}}^{\overrightarrow{h}}d\overrightarrow{x}\right)& =\frac{b}{2\pi it}\sum _{\alpha }\frac{1}{m_{\alpha }-1}{\left(\frac{d}{dz}\right)}^{m_{\alpha }-1}{e}^{z\log \left(\frac{a}{t}\right)}{|}_{z=-{\lambda }_{\alpha }}\\ & =\{\begin{array}{cc}\frac{b}{2\pi it}{\sum }_{\alpha }\frac{\log {\left(\frac{a}{t}\right)}^{m_{\alpha }-1}}{m_{\alpha }-1}{e}^{-{\lambda }_{\alpha }\log \left(\frac{a}{t}\right)},& \text{if}a>t,\\ 0,& \text{if}a\le t.\end{array}\end{array}$
When we now reverse the Laplace transform, we only have to integrate up to $a$.
$\begin{array}{cc}{\int }_0^1\exp (-na{\overrightarrow{x}}^{2\overrightarrow{k}})b{\overrightarrow{x}}^{\overrightarrow{h}}d\overrightarrow{x}& ={\int }_0^a\frac{b}{2\pi i}\sum _{\alpha }{a}^{-{\lambda }_{\alpha }}\frac{\log {\left(\frac{a}{t}\right)}^{m_{\alpha }-1}}{m_{\alpha }-1}{t}^{{\lambda }_{\alpha }-1}{e}^{-nt}dt|\tau =\frac{n}{a}t\\ & =\frac{b}{2\pi i}\sum _{\alpha }\frac{a^{-{\lambda }_{\alpha }}}{m_{\alpha }-1}{\int }_0^n\log {\left(\frac{n}{\tau }\right)}^{m_{\alpha }-1}{\left(\frac{a\tau }{n}\right)}^{{\lambda }_{\alpha }-1}{e}^{-\tau a}\frac{a}{n}d\tau \\ & =\frac{b}{2\pi i}\sum _{\alpha }\frac{n^{-{\lambda }_{\alpha }}}{m_{\alpha }-1}{\int }_0^n{\tau }^{{\lambda }_{\alpha }-1}{e}^{-\tau a}{\left(\log \right(n)-\log (\tau \left)\right)}^{m_{\alpha }-1}d\tau \\ & =\frac{b}{2\pi i}\sum _{\alpha }\frac{n^{-{\lambda }_{\alpha }}}{m_{\alpha }-1}{\int }_0^n{\tau }^{{\lambda }_{\alpha }-1}{e}^{-\tau a}\sum _{j=0}^{m_{\alpha }-1}(\frac{m_{\alpha }-1}{j})\log (n{)}^{m_{\alpha }-1-j}(-\log (\tau {)}^j)d\tau \\ & =\frac{b}{2\pi i}\sum _{\alpha }\frac{n^{-{\lambda }_{\alpha }}}{m_{\alpha }-1}\sum _{j=0}^{m_{\alpha }-1}(\frac{m_{\alpha }-1}{j})\log (n{)}^{m_{\alpha }-1-j}{\int }_0^n{\tau }^{{\lambda }_{\alpha }-1}{e}^{-\tau a}(-\log (\tau {)}^j)d\tau \end{array}$
This integral now is asymptotically going towards some number as $n$ goes to infinity. So we can pick out the term that is the highest order in $n$ as the one with the smallest ${\lambda }_{\alpha }$ and $j=0$.
$\begin{array}{c}{\int }_0^1\exp (-na{\overrightarrow{x}}^{2\overrightarrow{k}})b{\overrightarrow{x}}^{\overrightarrow{h}}d\overrightarrow{x}\propto n^{-{\lambda }_{\alpha \ast }}\log (n{)}^{m_{\alpha \ast }-1}+O(n^{-{\lambda }_{\alpha \ast }}\log \left(n{)}^{m_{\alpha \ast }-2}\right)\end{array}$
We can combine this into the Formula for Free Energy.
$\begin{array}{c}F({\omega }^{\ast },n)=n{K}_{min}+\lambda log\left(n\right)-(m-1)\log \left(\log \right(n\left)\right)\end{array}$
This expression tells us now, how much mass of the posterior distribution is centered around certain points in parameter space. The leading order term consists of the KL divergence between the training distribution and the distribution of our model. Intuitively, this corresponds to our training loss. In the second term the $\lambda$ inversely scales with the multiplicity of the minimum. Intuitively it corresponds to symmetries in the algorithm that the Network implements. The third term scales with the number of parameters that have a minimum of the same multiplicity. Intuitively this gives a set of parameters an extra advantage when multiple parameter sets implement algorithms that independently lead to minimal loss.

 

 

1. ^{^}

   I would have expected, that we integrate the whole vicinity:
   $\begin{array}{c}F({\omega }^{\ast },n)=-log\left({\int }_{-1}^1\exp \right(-n(K_{min}+a{\overrightarrow{x}}^{2\overrightarrow{k}})\left(b{\overrightarrow{x}}^{\overrightarrow{h}}d\overrightarrow{x}\right)\end{array}$
   This, however, makes the derivation quite ugly:
   $\begin{array}{cc}L\left({\int }_{-1}^1\exp \right(-na{\overrightarrow{x}}^{2\overrightarrow{k}}\left)b{\overrightarrow{x}}^{\overrightarrow{h}}d\overrightarrow{x}\right)& ={\int }_0^{\infty }\left({\int }_{-1}^1\exp \right(-na{\overrightarrow{x}}^{2\overrightarrow{k}}\left)b{\overrightarrow{x}}^{\overrightarrow{h}}d\overrightarrow{x}\right)e^{-nt}dt\\ & ={\int }_{-1}^1\delta (t-a{\overrightarrow{x}}^{2\overrightarrow{k}})b{\overrightarrow{x}}^{\overrightarrow{h}}d\overrightarrow{x}\\ M\left(L\right({\int }_{-1}^1\exp (-na{\overrightarrow{x}}^{2\overrightarrow{k}})b{\overrightarrow{x}}^{\overrightarrow{h}}d\overrightarrow{x}\left)\right)& ={\int }_0^{\infty }\left({\int }_{-1}^1\delta \right(t-a{\overrightarrow{x}}^{2\overrightarrow{k}}\left)b{\overrightarrow{x}}^{\overrightarrow{h}}d\overrightarrow{x}\right)t^{z}dt\\ & ={\int }_{-1}^1(a{\overrightarrow{x}}^{2\overrightarrow{k}}{)}^{z}b{\overrightarrow{x}}^{\overrightarrow{h}}d\overrightarrow{x}\\ & =a^{z}b\prod _i{\int }_{-1}^1{x}_i^{2z{k}_i+h_i}\\ & =\frac{a^{z}b}{{\prod }_i(2z{k}_i+h_i+1)}\prod _i(1^{2z{k}_i+h_i}-(-1{)}^{2z{k}_i+h_i})\end{array}$
   $\begin{array}{cc}L\left({\int }_{-1}^1\exp \right(-na{\overrightarrow{x}}^{2\overrightarrow{k}}\left)b{\overrightarrow{x}}^{\overrightarrow{h}}d\overrightarrow{x}\right)& =\frac{b}{2\pi i}\oint \frac{a^z{t}^{-z-1}}{{\prod }_i(2z{k}_i+h_i+1)}\prod _i(1-(-1{)}^{2z{k}_i+h_i})dz\\ & =\frac{b}{2\pi it}{\int }_{-i\infty }^{+i\infty }\frac{e^{z\log \left(\frac{a}{t}\right)}}{{\prod }_{\alpha }{\prod }_{i\in I_{\alpha }}2k_i(z+{\lambda }_i)}\prod _i(1-(-1{)}^{2k_i(z+{\lambda }_i)})dz\\ & =\frac{b}{2\pi it}\sum _{\alpha }\frac{1}{m_{\alpha }-1}{\left(\frac{d}{dz}\right)}^{m_{\alpha }-1}{e}^{z\log \left(\frac{a}{t}\right)}(1-(-1{)}^{2k_i(z+{\lambda }_i)}){|}_{z=-{\lambda }_{\alpha }}\\ & =\{\begin{array}{cc}\frac{b}{2\pi it}{\sum }_{\alpha }\frac{1}{m_{\alpha }-1}{e}^{-{\lambda }_{\alpha }\log \left(\frac{a}{t}\right)}{\sum }_j(\frac{m_{\alpha }-1}{j})\left(\log {\left(\frac{a}{t}\right)}^{m_{\alpha }-1-j}\right(2\pi i{k}_i{)}^j,& \text{if}a>t,\\ 0,& \text{if}a\le t.\end{array}\end{array}$
   We can proceed with the rest of the calculation with each summand $j$ and just have a smaller $m_{\alpha }$. This does not change anything about the leading terms, that are relevant for the Free Energy formula.

   I am still confused, about why the lecture went from 0 to 1, and would be grateful for an explanation in the comments.

6 comments

Comments sorted by top scores.