4.1 Naive stochastic gradient Hamiltonian Monte Carlo

A naive way is to directly substitute the noisy estimator in Equation 7 into the Hamiltonian dynamics in Equation 6:

In this case, the Hamiltonian is not guaranteed to be invariant:

When using a larger mini-batch size: , the variance is smaller: , resulting in At the limit, the total energy is preserved, which is the full-batch Hamiltonian Monte Carlo mentioned above.

When using a much smaller mini-batch size: , the noise induced by the mini-batch, , is large (e.g., in terms of matrix norm), resulting in Consequently, the Hamiltonian is no longer invariant.

To correct the error due to the effect of mini-batches, one needs to perform one Metropolis - Hastings step to either reject or accept the new state. Either running a short or long simulation (corresponding to a small or large in Algorithm Hamiltonian Monte Carlo), the cost of a Metropolis - Hastings step is still extremely large and wasteful if the sample is rejected. One workaround solution is to run a Metropolis - Hastings step on a subset of data instead of the entire dataset (Korattikara et al. 2014; Bardenet et al. 2014). There are, of course, some tradeoffs using such approaches.

Indeed, the joint distribution can be determined to be stationary or not by analysing the corresponding Fokker - Planck equation as shown in the Appendix about the stationary of stochastic gradient due to mini-batches. In this case, is proved to be non-stationary.

In (Chaudhari and Soatto 2018), the joint distribution in Equation 5 is assumed to be stationary under the stochastic dynamics in Equation 8. This is equivalent to proving that the left hand side term in the Fokker - Planck equation is zero: . The authors then analyse and show that the stationary distribution does not converge to the desired posterior distribution in general (Chaudhari and Soatto 2018). This is, however, only true if the stationary distribution exists. And in this case, we prove that it does not (the distribution is non-stationary as shown in Section stationary of stochastic gradient due to mini-batches).

4.2 Stochastic gradient Hamiltonian Monte Carlo with “friction”

One way to overcome the stochastic estimation for the gradient of the potential energy, , is to introduce a “friction” term to the momentum update: where: denotes friction coefficient matrix. One requirement for is that: (see the section on stationary SGD with injected noise for further details).

In this case, one can prove that the joint distribution is stationary.

To link this sampling to the stochastic gradient descent, one can sample and apply one leapfrog step as follows:

It can be simplified by substituting into the expression of to obtain: which has a similar form as the Stochastic Gradient Langevin Dynamics (Welling and Teh 2011).

2.1 Data generation

The data is generated as follows:

• draw the parameter from its prior: ,
• draw the parameter from its prior: ,
• draw a hidden sample from a prior distribution: , and
• draw an observable sample given as follows: ,

where and are the parameter of the model of interest.

NoteParameter

In many tutorials of EM, the parameter of the prior of the latent variable is often defined implicitly. In this post, it is defined explicitly to make the explanation easier to follow.

Such a data generation process is often visualised by the graphical model shown below

%%{
    init: {
        'theme': 'base',
        'themeVariables': {
            'primaryColor': '#ffffff'
        }
    }
}%%
flowchart LR
    subgraph data["data"]
        z((z)):::nonfilled-->x((x)):::filled;
    end
    pi((π)):::nonfilled-->z;
    theta((θ)):::nonfilled-->x;

    linkStyle default stroke: black;
    classDef nonfilled fill: none;
    style data fill: none;

2.2 Parameter inference

Given a set of observed i.i.d data , the general objective is to infer the posterior of the parameters and . Instead of inferring the exact posterior , which may be difficult in many cases, one can perform point estimate, such as MLE or maximise a posterior (MAP), which can be written as follows:

Due to the presence of the sum over the latent variable , the in-complete log-likelihood may not be evaluated directly on the joint distribution (especially when is continuous), making the optimisation difficult.

Fortunately, according to the data generation presented in Section 2.1, the completed log-likelihood can be evaluated easily:

Such an assumption allows EM to get around the difficulty when evaluating the expression in Equation 1.

The first two sub-steps combined are often known as the Expectation step (or E-step for short), while the last step is known as the Maximisation step (or M-step for short). These steps are then presented in the following sub-sub-sections.

2.3 Convergence of the EM algorithm

The following theorem proves that the EM algorithm improves the lower-bound after every iteration. For simplicity, the priors and are ignored from the proof below, but extending to include these prior terms is trivial.

3.1 Gaussian mixture models

The Gaussian mixture distribution can be written as a convex combination of Gaussian components: where: and .

%%{
    init: {
        'theme': 'base',
        'themeVariables': {
            'primaryColor': '#ffffff'
        }
    }
}%%
flowchart LR
    subgraph data["data"]
        direction LR
        z((z)):::rv --> x((x)):::rv
    end
    alpha((α)):::notfilled --> pi((π)):::params --> z
    sigma((Σ)):::params --> mu
    psi((Ψ)):::notfilled --> sigma
    nu((ν)):::notfilled --> sigma
    sigma --> x
    mu0((m)):::notfilled --> mu((μ)):::params --> x
    lambda((λ)):::notfilled --> mu

    style z fill: none
    classDef params stroke: #000, fill: none
    classDef rv stroke: #000
    classDef notfilled fill: none
    linkStyle default stroke: #000
    style data fill: none

3.2 Multinomial mixture models

Similar to the Gaussian mixture models, a multinomial mixture model can also be written as:

Note is given

Only the case where all the multinomial components have the same parameter (the number of trials) are considered. The reason is that optimising for an integer number is beyond the scope of this post.

%%{
    init: {
        'theme': 'base',
        'themeVariables': {
            'primaryColor': '#ffffff'
        }
    }
}%%
flowchart LR
    subgraph data["data"]
        direction LR
        z((z)):::rv --> x((x)):::rv
    end
    alpha((α)):::notfilled --> pi((π)):::params --> z;
    beta((β)):::params --> rho((ρ)):::params;
    rho --> x;

    style z fill: none
    classDef params stroke: #000, fill: none
    classDef rv stroke: #000
    classDef notfilled fill: none
    linkStyle default stroke: #000
    style data fill: none

3.2.3 Parameter inference with EM

E-step calculates the posterior of the latent variable given the data :

M-step In the M-step, we maximise the following expected completed log-likelihood w.r.t. and :

NoteProbability constrains on and

Due to the nature of a multinomial mixture model, both the parameters and are probability vectors.

The Lagrangian for can be written as: where is the Lagrange multiplier.

Taking derivative of the Lagrangian w.r.t. gives:

Setting the derivative to zero and solving for gives:

And since , one can substitute and find that . Thus:

Similarly, the Lagrangian of can be expressed as: where is the Lagrange multiplier. Taking derivative w.r.t. gives: Setting the derivative to zero and solving for gives: The constraint on as a probability vector leads to . Thus:

One can also refer to (Elmore and Wang 2003) for a similar derivation and result.

4.1 Binary classification

4.2 Multi-class classification

5.1 Some properties of Bregman divergence

This sub-section presents some properties of Bregman divergence, which can then be used in the bias-variance decomposition. Note that the notation and used in this section do not need to be label distribution, but can simply be the output of a model (without any normalization, e.g. no softmax). The case for label distributions will be considered in the subsequent section where the loss function is KL divergence.

5.2 Decomposition for Bregman divergence

The main result of bias-variance decomposition can be shown in the following:

1.1 Data generation model of a task

A data point of a task indexed by consists of an input and a corresponding label with . For simplicity, only two families of tasks – regression and classification – are considered in this thesis. As a result, the label is defined as for regression and as for classification, where is the number of classes.

Each data point in a task can be generated in 2 steps:

• generate the input by sampling from some probability distribution ,
• determine the label , where is the correct labelling function.

Both the probability distribution and the labelling function are unknown to the learning agent during training, and the aim of the supervised learning is to use the generated data to infer such labelling function .

For simplicity, we denote as the data generation model of task -th.

1.2 Task instance

To solve a task , one needs to obtain an optimal task-specific model , parameterised by , which minimises a loss function on the data of that task:

In practice, since both and are unknown, the data generation model is replaced by a dataset consisting of a finite number of data-points generated according to the data generation model , denoted as . The objective to solve that task is often known as empirical risk minimisation:

Since the loss function used is the same for each task family, e.g.  is NLL or variational-free energy, the subscript on the loss function is, therefore, dropped, and the loss is denoted as throughout this chapter. Furthermore, given the commonality of the loss function across all tasks, a task can, therefore, be simply represented by either its data generation model or the associated dataset .

1.3 Hyper-parameter optimisation

In single-task setting, the common way to tune or optimise a hyper-parameter is to split a given dataset into two disjoint subsets: where:

• is the training (or support) subset,
• is the validation (or query) subset.

Note that with this definition, , and and are not necessarily identical.

The subset is used to train the model parameter of interest , while the subset is used to validate the hyper-parameter, denoted by (we provide examples of the hyper-parameter in Section Formulation of meta-learning). Mathematically, hyper-parameter optimisation in the single-task setting can be written as the following bi-level optimisation:

We can extend the hyper-parameter optimisation from the two data subsets and to the general data generation model as the following: where and are the probability distributions of training and validation input data, respectively, and they are not necessarily identical.

Formulation of meta-learning

The setting of the meta-learning problem considered in this paper follows the task environment (Baxter 2000) that describes the unknown distribution over a family of tasks. Each task is sampled from this task environment and can be represented as , where and are the probability of training and validation input data, respectively, and are not necessarily identical. The aim of meta-learning is to use training tasks to train a meta-learning model that can be fine-tuned to perform well on an unseen task sampled from the same task environment.

Such meta-learning methods use meta-parameters to model the common latent structure of the task distribution . In this thesis, we consider meta-learning as an extension of hyper-parameter optimisation in single-task learning, where the hyper-parameter of interest — often called meta-parameter — is shared across many tasks. Similar to hyper-parameter optimisation presented in subsection hyper-parameter-optimisation, the objective of meta-learning is also a bi-level optimisation:

The difference between meta-learning and hyper-parameter optimisation is that the meta-parameter (also known as hyper-parameter) is shared across all tasks sampled from the task environment as highlighted in red colour in Equation 2.

In practice, the meta-parameter (or shared hyper-parameter) can be chosen as one of the followings:

• learning rate of gradient-based optimisation used to minimise the lower level objective function in Equation 2 to learn (Li et al. 2017),
• initialisation of model parameter (Finn et al. 2017),
• data representation or feature extractor (Vinyals et al. 2016; Snell et al. 2017),
• optimiser used to optimise the lower-level in Equation 2.

In this post, the meta-parameter is assumed to be the initialisation of model parameters. Formulation, derivation and analysis in the subsequent sections and chapters will, therefore, revolve around this assumption. Note that the analysis can be straight-forwardly extended to other types of meta-parameters with slight modifications.

In general, the objective function of meta-learning in Equation 2 can be solved by gradient-based optimisation, such as gradient descent. Due to the nature of the bi-level optimisation, the optimisation are often carried out in two steps. The first step is to adapt (or fine-tuned) the meta-parameter to the task-specific parameter . This corresponds to the optimisation in the lower-level, and can be written as: where is a hyper-parameter denoting the learning rate for task . For simplicity, the adaptation step in Equation 3} is carried out with only one gradient descent update.

The second step is to minimise the validation loss induced by the locally-optimal task-specific parameter evaluated on the validation subset w.r.t. the meta-parameter . This corresponds to the upper-level optimisation, and can be expressed as: where is another hyper-parameter representing the learning rate to learn .

The general algorithm of meta-learning using gradient-based optimisation is shown in Algorithm 1.

1.4 Second-order meta-learning

As shown in Equation 4, the optimisation for the meta-parameter requires the gradient of the validation loss averaged across tasks. Given that each task-specific parameter is a function of due to the lower-level optimisation in Equation 3, the gradient of interest can be expanded as: where the first equality is due to chain rule, and the second equality is the result that differentiates the gradient update in Equation 3. Note that in the second equality, we remove the transpose notation since the corresponding matrix is symmetric.

Thus, naively implementing such gradient would require to calculate the Hessian matrix $ $, resulting in an intractable procedure for large models, such as deep neural networks. To obtain a more efficient implementation, one can utilise the Hessian-vector product (Pearlmutter 1994) between the gradient vector and the Hessian matrix $ $ to efficiently calculate the gradient of the validation loss w.r.t. .

Another way to calculate the gradient of the validation loss w.r.t. the meta-parameter is to use implicit differentiation (Domke 2012; Rajeswaran et al. 2019; Lorraine et al. 2020). This approach is more advantaged since it does not need to stores the computational graph and takes gradient via chain rule. Such implicit differentiation technique reduces the memory usage and therefore, allows to work with large-scale models. However, the trade-off is the increasing computational time to apply the chain rule to calculate the gradient of interest.

Nevertheless, the implementations that compute the exact gradient of the validation loss w.r.t. without approximation are often referred to as second-order meta-learning.

1.5 First-order meta-learning

In practice, the Hessian matrix $ $ is often omitted from the calculation to simplify the update for the meta-parameter (Finn et al. 2017). The resulting gradient consists of only the gradient of validation loss , which is more efficient to calculate with a single forward-pass if auto differentiation is used. This approximation is often referred as first-order meta-learning, and the gradient of interest can be presented as:

REPTILE [Nichol et al. (2018)} — a variant first-order meta-learning — approximates further the gradient of validation loss by the difference , resulting in a much simpler approximation:

2.1 Fine-tuning

Fine-tuning is the most common technique in neural network based transfer learning (Pratt et al. 1991; Yosinski et al. 2014) where the last or a couple of last layers in a neural network pre-trained on a source task are replaced and fine-tuned on a target task. Formally, if is denoted as the forward function of the shared layers with shared parameters , where and are the parameters of the remaining layers specifically trained on source and target tasks, respectively, then the objective of fine-tuning can be expressed as:

where and are the data sampled from the source task and target task , respectively.

Although the objective of fine-tuning shown in Equation 5 is still a bi-level optimisation, it is easier to solve than the one in meta-learning due to the following reasons:

• The objective in fine-tuning has only one constrain corresponding to one source task, while meta-learning has several constrains corresponding to multiple training tasks.
• In fine-tuning, and are inferred separately, while in meta-learning, the task-specific parameter is a function of the meta-parameter, resulting in a more complicated correlation.

The downside of fine-tuning is the requirement of a reasonable number of training examples on the target task to fine-tune . In contrast, meta-learning leverages the knowledge extracted from several training tasks to quickly adapt to a new task with only a few training examples.

2.2 Domain adaptation and generalisation

Domain adaptation or domain-shift refers to the case when the joint data-label distribution on source and target are different, denoted as , or simply (Heckman 1979; Shimodaira 2000; Japkowicz and Stephen 2002; Daume III and Marcu 2006; Ben-David et al. 2007). The aim of domain adaptation is to leverage the model trained on source domain to available data in the target domain, so that the model adapted to the target domain can perform reasonably well. In other words, domain adaptation relies on a data transformation that produces a domain-invariant latent space. Mathematically, the transformation is obtained by minimising a divergence between the two transformed data distribution:

After obtaining the transformation , one can simply train a model using the transformed data of the source domain, and then use that model to make predictions on the target domain.

Given the optimisation in Equation 6, domain adaptation is different from meta-learning due to the following reasons:

• Domain adaptation assumes a shift in the task environments that generate source and target tasks, while meta-learning is based on the assumption of same task generation.
• Domain adaptation utilises information of data from target domain, while meta-learning does not have such access.

In general, meta-learning learns a shared prior or hyper-parameters to generalise for unseen tasks, while domain adaptation produces a model to solve a particular task in a specified target domain. Recently, there is a variance of domain adaptation, named domain generalisation, where the aim is to learn a domain-invariant model without any information of target domain. In this view, domain generalisation is very similar to meta-learning, and there are some works that employ meta-learning algorithms for domain generalisation (Li et al. 2018; Li et al. 2019).

2.3 Multi-task learning

Multi-task learning learns several related auxiliary tasks and a target task simultaneously to exploit the diversity of task representation to regularise and improve the performance on the target task (Caruana 1997). If the input is assumed to be the same across extra tasks and the target task , then the objective of multi-task learning can be expressed as: where and are the label, loss function and the classifier for task , respectively, and is the shared feature extractor for tasks.

Multi-task learning is often confused with meta-learning due to their similar nature extracting information from many tasks. However, the objective function of multi-task learning in Equation 7 is a single-level optimisation for the shared parameter and multiple task-specific classifier . It is, therefore, not as complicated as a bi-level optimisation seen in meta-learning as shown in Equation 2. Furthermore, multi-task learning aims to solve a number of specific tasks known during training (referred to as target tasks), while meta-learning targets the generalisation for unseen tasks in the future.

2.4 Continual learning

Continual or life-long learning refers to a situation where a learning agent has access to a continuous stream of tasks available over time, and the number of tasks to be learnt is not pre-defined (Chen and Liu 2018; Parisi et al. 2019). The aim is to accommodate the knowledge extracted from one-time observed tasks to accelerate the learning of new tasks without catastrophically forgetting old tasks (French 1999). In this sense, continual learning is very similar to meta-learning. However, continual learning most likely focuses on systematic design to acquire new knowledge in such a way that prevents interfering to the existing one, while meta-learning is more about algorithmic design to learn the new knowledge more efficiently. Thus, we cannot mathematically distinguish their differences as done in sub-sections Fine-tuning, Domain adaptation and generalisation and Multi-task learning . Nevertheless, continual learning criteria, especially catastrophic forgetting, can be encoded into meta-learning objective to advance further continual learning performance (Al-Shedivat et al. 2018; Nagabandi et al. 2019).

3.1 Mean square error in regression

In the regression:

• 
• is the identity function
• .

Hence, , resulting in: which agrees with the results in (Bishop and Nasrabadi 2006 - Eq.(5.84)).

3.2 Logistic regression

In this case:

• 
• is the sigmoid function
• .

The first derivative is expressed as:

The second derivative is therefore:

Hence: which agrees with the result derived in the literature (Bishop and Nasrabadi 2006 - Eq. (5.85)).

3.3 Cross entropy loss in classification

In this case:

• is the softmax function,
• .

According to the definition of the softmax function:

Hence, the derivative can be written as: and

An element of the Jacobian vector of the loss w.r.t. can be written as:

Hence, the Jacobian vector can be expressed as:

The Hessian matrix is given as:

Or, in the explicit matrix form:

1.1 Change of measure inequality for Kullback-Leibler divergence

1.2 Concentration inequality