# Probabilistic task modelling for meta-learning

Cuong Nguyen  Australian Institute for Machine Learning, University of
Adelaide, Australia  Thanh-Toan Do  Department of Data Science and AI, Monash
University, Australia  Gustavo Carneiro  Australian Institute for Machine
Learning, University of Adelaide, Australia

###### Abstract

We propose _probabilistic task modelling_ – a generative probabilistic model
for collections of tasks used in meta-learning. The proposed model combines
variational auto-encoding and latent Dirichlet allocation to model each task
as a mixture of Gaussian distribution in an embedding space. Such modelling
provides an explicit representation of a task through its task-theme mixture.
We present an efficient approximation inference technique based on variational
inference method for empirical Bayes parameter estimation. We perform
empirical evaluations to validate the _task uncertainty_ and _task distance_
produced by the proposed method through correlation diagrams of the prediction
accuracy on testing tasks. We also carry out experiments of task selection in
meta-learning to demonstrate how the task relatedness inferred from the
proposed model help to facilitate meta-learning algorithms.

##  1 Introduction

The latest developments in machine learning have enabled the field to solve
increasingly complex classification problems. Such complexity require high
capacity models, which in turn need a massive amount of annotated data for
training, resulting in an arduous, costly and even infeasible annotation
process. This has, therefore, motivated the research of novel learning
approaches, generally known as transfer learning, that exploit past experience
(in the form of models learned from other training tasks) to quickly learn a
new task using relatively small training sets.

Transfer-learning, and in particular, meta-learning, assumes the existent of a
task environment where training and testing tasks are i.i.d. sampled from the
same latent distribution. By modelling such environment through meta-
parameters that are shared across all tasks, meta-learning can solve an unseen
task more accurately and efficiently by learning how to solve many tasks
generated from the same distribution, even if each task contains a limited
number of training examples. Meta-learning has, therefore, progressed steadily
with many remarkable state-of-the-art results in several few-shot learning
benchmarks [36, 31, 10, 38, 28, 2]. However, current development of meta-
learning focuses on solving tasks without providing understanding on how tasks
are generated or correlated, potentially leading to sub-optimal solutions. In
fact, there is a large variation of prediction performance made by various
meta-learning algorithms reported in [9, Figure 1] or shown in Figure 1,
suggesting that not all testing tasks are equally related to training tasks.
This motivates our work to model and represent tasks in a latent task-theme
space. The new task representation allows further downstream works, such as
task similarity or active task selection, to be developed to gain insights
into, or even improve, the prediction performance of different meta-learning
algorithms.

In this paper, we propose probabilistic task modelling (PTM) – a graphical
model that combines variational auto-encoding (VAE) [16] and Gaussian latent
Dirichlet allocation (LDA) [8] – to model tasks used in meta-learning111This
is a comprehensive work extended from our workshop paper [22]. Note that PTM
itself is not a meta-learning method. With this modelling approach, the
dataset associated with each task can be modelled as a mixture of finite
Gaussian distributions, allowing to represent tasks in a latent task-theme
simplex via its mixture coefficient vector. Such representation provides a
convenient way to quantitatively measure task uncertainty or task relatedness,
which can be utilised in active task selection for meta-learning.

Figure 1: The results locally produced for MAML on 15,504 available 5-way
1-shot mini-ImageNet testing tasks vary from 20 to 70 percent accuracy,
suggesting that not all testing tasks are equally related to training tasks.

##  2 Related Work

The proposed approach is closely related to Task2Vec [1] when modelling tasks
for meta-learning. In Task2Vec, a task is represented by an embedding computed
from the Fisher information matrix associated with the task-specific
classifier. In PTM, a task is represented by the variational distribution of
task-theme mixture, which is a part of the graphical model describing the task
generation process. The two methods, therefore, differ at the modelling
mechanism: Task2Vec follows a deterministic approach, while PTM is a
probabilistic method. Such difference provides an advantage of PTM over
Task2Vec, which includes modelling uncertainty into the task representation.
In addition, PTM is more efficient than Task2Vec at inference when predicting
task representation, since PTM only needs a single forward pass, while
Task2Vec requires to re-train or fine-tune the task-specific classifier and
calculate the Fisher information matrix for the task that needs to be
presented. There is also a slightly similar work recently published about task
modelling [15].

Our work is related to task similarity estimation, which has been intensively
studied in the field of multi-task learning. Some remarkable examples in this
area include task-clustering using k-nearest neighbours [33], or modelling
common prior between tasks as a mixture of distributions [3, 37]. Another
approach is to formulate multi-task learning as a convex optimisation problem
either to cluster tasks and utilise the clustering results to fast track the
learning [14], or to learn task relationship through task covariance matrices
[40]. Other approaches provided theoretical guarantees when learning the
similarity or relationship between tasks [30]. Recently, the _taskonomy_
project [39] was conducted to carry out extensive experiments on 26 computer-
vision tasks to empirically analyse the correlation between those tasks. Other
works [34, 24] take a slightly different approach by investigating the
correlation of the label distributions between the tasks of interest to
measure task similarity. One commonality among all studies above is their
reliance on task-specific classifiers which are used to quantify task
relatedness. In contrast, our proposed method explicitly models tasks without
the help of any task-specific classifier, making it more efficient in training
and prediction.

Our work is also connected to finite mixture models [25], such as the _latent
Dirichlet allocation_ [5], which analyses and summarises text data in topic
modelling, or categorises natural scenes in computer vision [18]. LDA assumes
that each document within a given corpus can be represented as a mixture of
finite categorical distributions, where each categorical distribution is a
latent topic shared across all documents. Training an LDA model or its
variants on a large text corpus is challenging, so several approximate
inference techniques have been proposed, ranging from mean-field variational
inference (VI) [5], collapsed Gibbs’ sampling [12] and collapsed VI [32].
Furthermore, several online inference methods have been developed to increase
the training efficiency for large corpora [6, 13, 11]. Our work is slightly
different from the modelling of the conventional LDA, where we do not use the
data directly, but embed it into a latent space. In short, our approach is a
combination of VAE [16] and LDA to model the dataset associated with a task.
Our approach considers word as continuous data, instead of the discrete data
represented by a bag-of-word vector generally used by LDA-based topic
modelling methods. The resultant model in the embedding latent space is,
therefore, similar to the Gaussian LDA [8] for word embedding in topic
modelling.

##  3 Probabilistic task modelling

To relate task modelling to topic modelling, we consider a task as a document,
and a data point as a word. Given these analogies, we can use LDA [5] – a
popular topic model – to model tasks for meta-learning. However, simply
applying LDA for task modelling would not scale well with high-dimensional
data and large datasets. We, therefore, propose to employ the VAE [16] to
reduce the dimension of the data, and use the inferred embeddings of data as
words to model tasks. Due to the nature of VAE, the latent variables are often
continuous, not discrete as the bag-of-words used in the conventional LDA. We,
therefore, replace the _categorical_ word-topic distributions in LDA by
_Gaussian_ task-theme222Task-theme is inspired by [18] distributions. Given
these assumptions, each task can be modelled as a mixture of $K$ Gaussian
task-themes, allowing to represent tasks by their inferred _task-theme
mixture_ vectors in the latent task-theme simplex as illustrated in Figure 2.
Hence, it is beneficial to utilise this representation for further downstream
tasks, such as measuring task similarity.

The graphical model of the proposed PTM is shown in Figure 3, where there are
$T$ tasks, and each task consists of $N$ data points, denoted as $\mathbf{x}$.
To simplify the formulation and analysis, $N$ is assumed to be fixed across
all tasks, but the extension of varying $N$ is trivial and can be implemented
straightforwardly. Under these assumptions, a task can be generated as
follows:

  * •

Initialise the Dirichlet prior for task-theme mixture:
$\\{\alpha_{k}\\}_{k=1}^{K}$, where $\alpha\in\mathbb{R}_{+}$

  * •

Initialise means and covariance matrices of $K$ Gaussian task-themes
$\\{\bm{\mu}_{k},\bm{\Sigma}_{k}\\}_{k=1}^{K}$, where
$\bm{\mu}_{k}\in\mathbb{R}^{D}$, $\bm{\Sigma}_{k}\in\mathbb{R}^{D\times D}$ is
positive definite matrix, and $D$ is the dimension of the data embedding

  * •

For task $\mathcal{T}_{i}$ in the collection of $T$ tasks:

    * –

Choose a task-theme mixture vector: $\bm{\pi}_{i}\leavevmode\nobreak\
\sim\leavevmode\nobreak\ \mathrm{Dirichlet}\left(\bm{\pi};\bm{\alpha}\right)$

    * –

For data point $n$-th of task $\mathcal{T}_{i}$:

      * *

Choose an task-theme assignment one-hot vector:
$\mathbf{z}_{in}\leavevmode\nobreak\ \sim\leavevmode\nobreak\
\mathrm{Categorical}\left(\mathbf{z};\bm{\pi}_{i}\right)$

      * *

Choose an embedding of the data point:
${\mathbf{u}_{in}\sim\mathcal{N}\left(\mathbf{u};\bm{\mu}_{k},\bm{\Sigma}_{k}\right)}$,
where: $z_{ink}=1$

      * *

Generate the data point from a decoder $h$ parameterised by $\theta$:
${\mathbf{x}_{in}=h(\mathbf{u}_{in};\theta)}$.

Figure 2: An example of a task-theme simplex where each task is represented by
a 3-dimensional mixture vector.

To model tasks according to the task generation described above, we need to
infer the task-agnostic (or meta) parameters of interest, namely
$\bm{\mu},\bm{\Sigma},\bm{\alpha}$ and $\theta$. However, due to the
complexity of the graphical model shown in Figure 3, the exact inference for
the posterior $p(\bm{\mu},\bm{\Sigma},\bm{\alpha},\theta|\mathbf{x})$ is
intractable, and therefore, the estimation must rely on approximate inference.
For simplicity, maximum likelihood estimation (MLE) is used as the objective
function:

$\max_{\bm{\mu},\bm{\Sigma},\bm{\alpha},\theta}\ln
p(\mathbf{x}|\bm{\mu},\bm{\Sigma},\bm{\alpha},\theta).$ (1)

Although MLE simplifies the learning for the meta-parameters of interest, the
log-likelihood in (1) is still difficult to evaluate for the optimisation. One
workaround solution is to find its lower-bound, and maximise the lower-bound
instead of maximising the log-likelihood itself. This approach is analogous to
the variational inference, which has been widely used to infer the latent
parameters of VAE and LDA models.

Figure 3: The graphical model used in task modelling. The solid arrows denote
data generation, while the dashed arrows stand for inference. The boxes are
plates representing replicates. The shading nodes denote observable variables,
while the white nodes denote latent variables.

Since the proposed PTM is a combination of VAE and LDA, the derivation for the
lower-bound of the likelihood in (1) can be divided into 2 steps, where the
first step is analogous to the lower bound of a VAE, and the second step is
similar to the plain LDA model.

In the first step, the latent variable $\mathbf{u}$ is introduced, so that the
log-likelihood $\ln p(\mathbf{x}|\bm{\mu},\bm{\Sigma},\bm{\alpha},\theta)$ can
be bounded below by Jensen’s inequality:

$\ln
p(\mathbf{x}|\bm{\mu},\bm{\Sigma},\bm{\alpha},\theta)\geq\mathsf{L}_{\mathrm{VAE}},$
(2)

where the lower-bound is defined as:

$\displaystyle\mathsf{L}_{\mathrm{VAE}}=\mathbb{E}_{q(\mathbf{u})}\left[\ln
p\left(\mathbf{x}|\mathbf{u},\theta\right)+\ln
p\left(\mathbf{u}|\bm{\mu},\bm{\Sigma},\bm{\alpha}\right)-\ln
q(\mathbf{u})\right],$ (3)

with $q(\mathbf{u})$ being the variational distribution for the latent
variable $\mathbf{u}$.

Following the conventional VI for VAE [16], the variational distribution for
the latent embedding $\mathbf{u}$ is assumed to be a Gaussian distribution
with diagonal covariance matrix:

$q(\mathbf{u})=\mathcal{N}\left(\mathbf{u};\mathbf{m},\mathrm{diag}\left(\mathbf{s}^{2}\right)\right).$
(4)

In addition, the parameters $\mathbf{m}$ and $\mathbf{s}$, which represent the
distribution encoding the data $\mathbf{x}$, are modelled by a neural network
(also known as an encoder) $f$ parameterised by $\phi$:

$\begin{bmatrix}\mathbf{m}^{\top}&\mathbf{s}^{\top}\end{bmatrix}^{\top}=f(\mathbf{x};\phi).$
(5)

Hence, instead of maximising the marginal log-likelihood in (1), the lower-
bound in (3) is maximised, resulted in the alternative objective:

$\max_{\bm{\mu},\bm{\Sigma},\bm{\alpha},\theta}\max_{\phi}\mathsf{L}_{\mathrm{VAE}}.$
(6)

One difficulty in maximising the lower-bound in (6) is the evaluation for the
prior $\ln p\left(\mathbf{u}|\bm{\mu},\bm{\Sigma},\bm{\alpha}\right)$ of the
latent embedding $\mathbf{u}$ shown in Eq. (3). In vanilla VAE, the embedding
prior is often modelled as some standard distributions, such as Gaussian or
Beta, resulting in a tractable solution. In this paper, the prior is modelled
as a Gaussian mixture model, making the solution intractable. However, since
this prior is the marginal log-likelihood in the conventional LDA model, we
can apply techniques developed for LDA methods to approximate this term. Here,
we employ the VI approach in which the term is bounded below by Jensen’s
inequality:

$\ln
p\left(\mathbf{u}|\bm{\mu},\bm{\Sigma},\bm{\alpha}\right)\geq\mathsf{L}_{\mathrm{LDA}}(\mathbf{u},q(\mathbf{z},\bm{\pi})),$
(7)

where:

$\displaystyle\mathsf{L}_{\mathrm{LDA}}\left(\mathbf{u},q(\mathbf{z},\bm{\pi})\right)=\mathbb{E}_{q(\mathbf{z},\bm{\pi})}\left[\ln
p(\mathbf{u}|\mathbf{z},\bm{\mu},\bm{\Sigma})+\ln p(\mathbf{z}|\bm{\pi})+\ln
p(\bm{\pi}|\bm{\alpha})-\ln q(\mathbf{z})-\ln q(\bm{\pi})\right],$ (8)

with $q(\mathbf{z},\bm{\pi})$ being the variational distribution for
$\mathbf{z}$ and $\bm{\pi}$. This corresponds to the second step in the
derivation.

Similar to LDA [5], the variational distribution $q(\mathbf{z},\bm{\pi})$ is
assumed to be fully factorised and followed the conjugate priors:

$q(\mathbf{z},\bm{\pi})=\prod_{i=1}^{T}q(\bm{\pi}_{i};\bm{\gamma}_{i})\prod_{n=1}^{N}q(\mathbf{z}_{in};\mathbf{r}_{in}),$
(9)

where:

$\displaystyle q(\bm{\pi}_{i};\bm{\gamma}_{i})$
$\displaystyle=\mathrm{Dirichlet}\left(\bm{\pi}_{i};\bm{\gamma}_{i}\right)$
(10) $\displaystyle q(\mathbf{z}_{in};\mathbf{r}_{in})$
$\displaystyle=\mathrm{Categorical}\left(\mathbf{z}_{in};\mathbf{r}_{in}\right),$
(11)

with $\mathbf{r}$ and $\bm{\gamma}$ being the parameters of the variational
distribution $q(\mathbf{z},\bm{\pi})$.

In practice, $q(\mathbf{z},\bm{\pi})$ is obtained as the maximiser of the
lower-bound
$\mathsf{L}_{\mathrm{LDA}}\left(\mathbf{u},q(\mathbf{z},\bm{\pi})\right)$ on
the embedding data $\mathbf{u}$. A naive implementation to infer such
variational distribution is problematic since the data is used twice: one to
optimise $q(\mathbf{z},\bm{\pi})$, and the other is to optimise for
$q(\mathbf{u})$ in (6), which may result in overfitting. To avoid this issue,
we employ the _empirical Bayes_ approach by splitting the training dataset of
each task into two halves, where one half of data in a task, denoted as
$\mathbf{u}^{(t)}$, is used to obtain $q(\mathbf{z},\bm{\pi})$, while the
other half, denoted as $\mathbf{u}^{(v)}$, is used for the optimisation in
(6). This approach is analogous to the _empirical Bayes_ meta-learning [10,
23], where one part of data is used for task-adaptation (often known as inner-
loop), while the other part is used to learn the meta-parameter (often known
as outer-loop).

Given this modelling approach, the objective function can be formally written
as a bi-level optimisation:

$\displaystyle\max_{\bm{\mu},\bm{\Sigma},\bm{\alpha},\theta,\phi}$
$\displaystyle\mathsf{L}\left(\mathbf{u}^{(v)},q^{*}\left(\mathbf{z},\bm{\pi}\right)\right)$
(12) s.t.: $\displaystyle
q^{*}\left(\mathbf{z},\bm{\pi}\right)=\arg\max_{q(\mathbf{z},\bm{\pi})}\mathbb{E}_{q(\mathbf{u}^{(t)};\phi)}\left[\mathsf{L}_{\mathrm{LDA}}\left(\mathbf{u}^{(t)},q(\mathbf{z},\bm{\pi})\right)\right],$

where

$\begin{split}&\mathsf{L}\left(\mathbf{u}^{(v)},q^{*}\left(\mathbf{z},\bm{\pi}\right)\right)=\mathbb{E}_{q\left(\mathbf{u}^{(v)};\phi\right)}\left[\mathsf{L}_{\mathrm{LDA}}\left(\mathbf{u}^{(v)},q^{*}(\mathbf{z},\bm{\pi})\right)+\ln
p\left(\mathbf{x}^{(v)}|\mathbf{u}^{(v)},\theta\right)-\ln
q\left(\mathbf{u}^{(v)};\phi\right)\right].\end{split}$ (13)

Due to the assumptions made in Eqs. 4, 10 and 11, prior conjugate can be
applied to simplify the evaluation for all the terms in (8) w.r.t. the
variational distribution $q(.)$. Details of the evaluation can be referred to
Appendix A. In addition, the optimisation for the meta-parameters in (12) is
based on gradient ascent, and carried out in two steps, resulting in a process
analogous to the expectation-maximisation (EM) algorithm. In the E-step
(corresponding to the optimisation for the lower-level in (12)), the task-
specific variational-parameters $\mathbf{r}$ and $\bm{\gamma}$ are iteratively
updated, while holding the meta-parameters
$\bm{\mu},\bm{\Sigma},\bm{\alpha},\theta$ and $\phi$ fixed. In the M-step
(corresponding to the optimisation for the upper-level), the meta-parameters
are updated using the values of the task-specific variational-parameters
obtained in the E-step. Note that the inference for the task-theme parameters
$\bm{\mu}$ and $\bm{\Sigma}$ are similar to the estimation of Gaussian mixture
model [4, Chapter 9]. Please refer to Appendix B for more details on the
optimisation.

Conventionally, the iterative updates in the E-step and M-step require a full
pass through the entire collection of tasks. This is, however, very slow and
even infeasible since the number of tasks, $T$, is often in the magnitude of
millions. We, therefore, propose an online VI inspired by the online learning
for LDA [13] to make the solution more scalable when inferring the meta-
parameters. For each task $\mathcal{T}_{i}$, we perform the EM to obtain the
task-specific parameters (denoted by a tilde on top of variables) that are
locally optimal for that task. The meta parameters of interest are then
updated as a weighted average between their previous values and the task-
specific values:

$\displaystyle\bm{\mu}$
$\displaystyle\leftarrow(1-\rho_{i})\bm{\mu}+\rho_{i}\tilde{\bm{\mu}}$ (14)
$\displaystyle\bm{\Sigma}$
$\displaystyle\leftarrow(1-\rho_{i})\bm{\Sigma}+\rho_{i}\tilde{\bm{\Sigma}}$
$\displaystyle\bm{\alpha}$
$\displaystyle\leftarrow\bm{\alpha}-\rho_{i}\underbrace{\tilde{\bm{\alpha}}_{i}}_{\mathbf{H}^{-1}\mathbf{g}},$

where $\rho_{i}=(\tau_{0}+i)^{-\tau_{1}}$ with $\tau_{0}\geq 0$ and
$\tau_{1}\in(0.5,1]$ [13], $\mathbf{g}$ is the gradient of
$\mathsf{L}_{\mathrm{LDA}}$ w.r.t. $\bm{\alpha}$, and $\mathbf{H}$ is the
Hessian matrix. The learning for the encoder $\phi$ and the decoder $\theta$
follows the conventional learning by stochastic gradient ascent. The complete
learning algorithm for the proposed probabilistic task modelling is shown in
Algorithm 1.

Also, instead of updating the meta-parameters as in (14) when observing a
single task, we use multiple or a mini-batch of tasks to reduce the effect of
measurement noise. The mini-batch version requires a slight modification in
the formulation presented above, where we calculate the average of all task-
specific parameters for tasks in the same mini-batch, and use that as the
task-specific value to update the corresponding meta parameters.

Algorithm 1 Online probabilistic task modelling

1:procedure Training

2: Initialise LDA parameters:
$\\{\bm{\mu}_{k},\bm{\Sigma}_{k},\alpha_{k}\\}_{k=1}^{K}$

3: Initialise encoder $\phi$ and decoder $\theta$

4: for each mini-batch of $T_{\mathrm{mini}}$ tasks do

5: for $i=1:T_{\mathrm{mini}}$ do

6: Split data into $\\{\mathbf{x}_{i}^{(t)},\mathbf{y}_{i}^{(t)}\\}$ and
$\\{\mathbf{x}_{i}^{(v)},\mathbf{y}_{i}^{(v)}\\}$

7: $\mathbf{m}^{(t)}_{i},\mathbf{s}_{i}^{(t)}\leftarrow
f(\mathbf{x}_{i}^{(t)};\phi)$

8: $\mathbf{m}^{(v)}_{i},\mathbf{s}_{i}^{(v)}\leftarrow
f(\mathbf{x}_{i}^{(v)};\phi)$

9: $\bm{\gamma},\mathbf{r}\leftarrow$
E-step($\mathcal{N}(\mathbf{u};\mathbf{m}_{i}^{(t)},\mathrm{diag}((\mathbf{s}_{i}^{(t)})^{2}))$)

10: Calculate
$\mathsf{L}\left(\mathbf{u}_{i}^{(v)},q^{*}_{i}\left(\mathbf{z},\bm{\pi}\right)\right)$
$\triangleright$ Eq. (13)

11: Calculate local task-themes
$\tilde{\bm{\mu}}_{i},\tilde{\bm{\Sigma}}_{i},\tilde{\bm{\alpha}}_{i}$

12: end for

13:
$\overline{\mathsf{L}}=\frac{1}{T}\sum_{i=1}^{T}\mathsf{L}\left(\mathbf{u}_{i}^{(v)},q^{*}_{i}\left(\mathbf{z},\bm{\pi}\right)\right)$

14:
$\bm{\mu},\bm{\Sigma},\bm{\alpha}\leftarrow\mathrm{online\\_LDA}\left(\tilde{\bm{\mu}}_{1:T},\tilde{\bm{\Sigma}}_{1:T},\tilde{\bm{\alpha}}_{1:T}\right)$

15: $\theta,\phi\leftarrow\mathrm{SGD}\left(-\overline{\mathsf{L}}\right)$
$\triangleright$ gradient ascent

16: end for

17: return $\bm{\mu},\bm{\Sigma},\bm{\alpha},\theta,\phi$

18:end procedure

19:

20:procedure
E-step($\mathcal{N}(\mathbf{u};\mathbf{m},\mathrm{diag}(\mathbf{s}^{2}))$)

21: Initialise $\mathbf{r},\bm{\gamma}$

22: repeat

23: calculate the un-normalised $r_{ink}$ $\triangleright$ Eq. (25)

24: normalise $\mathbf{r}_{in}$ such that $\sum_{k=1}^{K}r_{ink}=1$

25: calculate $\gamma_{ik}$ $\triangleright$ Eq. (29)

26: until $\frac{1}{K}\leavevmode\nobreak\ |\text{change in
}\bm{\gamma}|<\text{threshold}$

27: return $\bm{\gamma},\mathbf{r}$

28:end procedure

Although the reconstruction term $\ln
p(\mathbf{x}^{(v)}|\mathbf{u}^{(v)},\theta)$ in (12) is used to model the
likelihood of un-labelled data, it can straightforwardly be extended to a
labelled data pair $\\{\mathbf{x}^{(v)},\mathbf{y}^{(v)}\\}$ by introducing
the parameter $\mathbf{w}$ of a classifier. In that case, the reconstruction
term can be expressed as:

$\displaystyle\ln
p(\mathbf{x}^{(v)},\mathbf{y}^{(v)}|\mathbf{u}^{(v)},\theta,\mathbf{w})=\underbrace{\ln
p(\mathbf{y}^{(v)}|\mathbf{u}^{(v)},\mathbf{w})}_{\text{negative
classification loss}}+\underbrace{\ln
p(\mathbf{x}^{(v)}|\mathbf{u}^{(v)},\theta)}_{\text{negative reconstruction
loss}}.$ (15)

In general, $\mathbf{w}$ can be either a task-specific parameter generated
from an additional meta-parameter shared across all tasks – corresponding to
_empirical Bayes_ meta-learning (e.g. using train-test split to learn hyper-
parameters) algorithms [10, 23], or a meta-parameter itself – corresponding to
_metric_ meta-learning [36, 31]. For simplicity, we will use the latter
approach relying on the prototypical network [31] with Euclidean distance on
the data embedding $\mathbf{u}$, to calculate the classification loss on
labelled data. This reduces the need to introduce an additional parameter
$\mathbf{w}$ into our modelling.

### Task representation

Given the inferred meta-parameters, including the task-themes
$\\{\bm{\mu}_{k},\bm{\Sigma}_{k}\\}_{k=1}^{K}$, the Dirichlet prior
$\\{\bm{\alpha}_{l}\\}_{l=1}^{L}$, the encoder $\phi$ and the decoder
$\theta$, we can embed the data of a task into a latent space, and calculate
its variational Dirichlet posterior of the task-theme _mixing coefficients_
$q(\bm{\pi};\bm{\gamma}_{i})$. The obtained distribution can be used represent
the corresponding task in the latent task-theme simplex as illustrated in
Figure 2. This new representation of tasks has two advantages comparing to the
recently proposed task representation Task2Vec [1]: (i) it explicitly models
and represents tasks without the need of any pre-trained networks to use as a
probe network, and (ii) it uses a probability distribution, instead of a
vector as in Task2Vec, allowing to include modelling uncertainty when
representing tasks. Given the probabilistic nature of PTM, we can use the
entropy of the inferred task-theme mixture distribution
$q(\bm{\pi};\bm{\gamma}_{i})$ as a measure of _task uncertainty_. In section
4.1, we empirically show that this measure correlates to the generalisation or
test performance.

In addition, the representation produced by PTM can be used to quantitatively
analyse the similarity or _distance_ between two tasks $i$ and $j$ through a
divergence between $q(\bm{\pi};\bm{\gamma}_{i})$ and
$q(\bm{\pi};\bm{\gamma}_{j})$. Commonly, symmetric distances, such as Jensen-
Shannon divergence, Hellinger distance, or earth’s mover distance are employed
to calculate the divergence between distributions. However, it is argued that
similarity should be represented as an asymmetric measure [35]. This is
reasonable in the context of transfer learning, since knowledge gained from
learning a difficult task might significantly facilitate the learning of an
easy task, but the reverse might not always have the same level of
effectiveness. In light of asymmetric distance, we decide to use Kullback-
Leibler (KL) divergence, denoted as $D_{\mathrm{KL}}[.\|.]$, to measure _task
distance_. As $D_{\mathrm{KL}}\left[P\|Q\right]$ is defined as the information
lost when using a code optimised for $Q$ to encode the samples of $P$, we,
therefore, calculate
$D_{\mathrm{KL}}\left[q(\bm{\pi};\bm{\gamma}_{T+1})\|q(\bm{\pi};\bm{\gamma}_{i})\right]$,
where $i\in\\{1,\ldots,T\\}$, to assess how the training task
$\mathcal{T}_{i}$ differs from the learning of the novel task
$\mathcal{T}_{T+1}$.

##  4 Experiments

In this section, we empirically validate the two properties of PTM – _task
uncertainty_ and _task distance_ – through task distance matrix and
correlation diagrams. We also show two applications of the proposed approach
used in active task selection for inductive and transductive life-long meta-
learning. The experiments are based on the $n$-way $k$-shot tasks formed from
Omniglot [17] and mini-ImageNet [36] – the two widely used datasets to
evaluate the performance of meta-learning algorithms.

The Omniglot dataset consists of 1623 different handwritten characters from 50
different alphabets, where each character was drawn in black and white by 20
different people. Instead of using random train-test split that mixes all
characters, the original split [17] is used to yield finer-grained
classification tasks. In addition to the task forming based on randomly mixing
characters of many alphabets, the two-level hierarchy of alphabets and
characters are utilised to increase the difficulty of the character
classification. Note that no data augmentations, such as rotating images by
multiples of 90 degrees, is used throughout the experiments. Also, all images
are down-sampled to 64-by-64 pixel2 to simplify the image reconstruct in the
decoder.

The mini-ImageNet dataset comprises a small version of ImageNet, which
contains 100 classes taken from ImageNet, and each class has 600 colour
images. We follow the common train-test split that uses 64 classes for
training, 16 classes for validation, and 20 classes for testing [27]. Similar
to Omniglot, all images are also in 64-by-64 pixel2.

The encoder used in the experiments consists of 4 convolutional modules, where
each module has a convolutional layer with 4-by-4 filters and 2-by-2 stride,
followed by a batch normalisation and a leaky rectified linear activation
function with a slope of 0.01. The output of the last convolutional layer is
flattened and connected to a fully connected layer to output the desired
dimension for the latent variable $\mathbf{u}$. The decoder is designed
similarly, except that the convolutional operator is replaced by the
corresponding transposed convolution. For the Omniglot dataset, the number of
filters within each convolutional layer of the encoder is 8, 16, 32, and 64,
respectively, and the dimension of $\mathbf{u}$ is 64. For mini-ImageNet
dataset, these numbers are 32, 64, 128 and 256, and the dimension of
$\mathbf{u}$ is 128. The reconstruction loss follows the negative log-
likelihood of the continuous Bernoulli distribution [19], which is often known
as binary cross-entropy, while the classification loss is based on the
prototypical network used in metric learning. The training subset of each
task, $\mathbf{u}^{(t)}_{i}$, is used to calculate the class prototypes, and
the classification loss is based on the soft-max function of the distances
between the encoding of each input image to those prototypes [31]. The
optimiser used is Adam with the step size of $2\times 10^{-4}$ to optimise the
parameters of the encoder and decoder after every mini-batch consisting of 20
tasks. For the LDA part, a total of $K=8$ task-themes is used. The Dirichlet
prior is assumed to be symmetric with a concentration $\alpha=1.1$ across both
datasets. The parameters of the learning rate used in the online LDA are
$\rho_{0}=10^{6}$ and $\rho_{1}=0.5$. A total of $10^{6}$ episodes are used to
train PTM on both datasets. We note that setting $\alpha>1$ enforces every
task to be modelled as a mixture of many task-themes, avoiding the task-themes
collapsing into a single task-theme during training. The phenomenon of task-
theme collapse when $\alpha<1$ is not observed in LDA, but in PTM due to the
integration of VAE. At the beginning of training, the encoder is inadequate,
producing mediocre embedding features. The resulting features, combined with
$\alpha<1$, makes a task more likely to be represented by a single task-theme.
By learning solely from that task-theme, the encoder is pushed to bias further
toward to that task-theme, making only one task-theme distribution updated,
while leaving others unchanged. When $\alpha>1$, all the task-themes
contribute to the representation of a task, so they can be learnt along with
the encoder. Please refer to
https://github.com/cnguyen10/probabilistic_task_modelling for the implemented
code.

Figure 4: The matrix of log KL distances between Omniglot tasks shows that
tasks that are generated from the same alphabet are closer together, denoted
as the dark green blocks along the diagonal. The matrix is asymmetric due to
the asymmetry of the KL divergence used as the task distance.

(a)

(b)

(c)

(d)

Figure 5: Correlation diagrams between prediction accuracy made by MAML on 100
5-way 1-shot testing tasks versus: 5(a) and 5(b) entropy of the inferred task-
theme mixture distributions, and 5(c) and 5(d) the KL distances from testing
to training tasks. The results show that largest the task entropy or
distances, the worse the testing performance. The blue dots are the prediction
made the MAML and PTM, the solid line is the mean of Bayesian Ridge
regression, and the shaded areas correspond to $\pm 1$ standard deviation
around the mean.

###  4.1 Task distance matrix and correlation diagrams

Task distance matrix is used as one of the tools to qualitatively validate the
prediction made by PTM. In particular, the hypothesis is that the PTM would
predict small distances for tasks that are close together. Since the labels
specifying the closedness of tasks are unknown, we utilise the hierarchical
structure of Omniglot dataset to form tasks. Each task is generated by firstly
sampling an alphabet, and then choosing characters in that alphabet. Under
this strategy, tasks formed from the same alphabet would have small distances
comparing to tasks from different alphabets. Figure 4 shows the task distances
between 50 testing tasks of Omniglot dataset, where each block of 5 tasks on
rows and columns of the task distance matrix corresponds to a group of tasks
sampled from the same alphabet. The result, especially the square 5-task-
by-5-task blocks along the diagonal, agrees well with the hypothesis. Note
that the distance matrix shown in Figure 4 is asymmetric due to the asymmetric
nature of the KL divergence used to measure task distance.

We use a correlation diagram between prediction accuracy made by MAML and the
_task entropy_ produced by PTM as another verification. Since the _task
entropy_ denotes the uncertainty when modelling a task, we hypothesise that it
proportionally relates to the difficulty when learning that task. To construct
the correlation diagram, we firstly train a meta-learning model based on MAML
using the training tasks of the two datasets, and evaluating the performance
on 100 random testing tasks. Secondly, we calculate the _task entropy_ for
those 100 testing tasks. Finally, we plot the prediction accuracy and task
entropy in Figures 5(a) and 5(b). The results on both datasets show that the
higher the task uncertainty, the worse the test performance. This observation,
therefore, agrees with our hypothesis about _task entropy_.

We conduct another correlation diagram between training-testing task distance
and the test performance to verify further the proposed PTM. Our hypothesis is
the inverse proportion between training-testing task distance and prediction
accuracy. A similar experiment as in task uncertainty is carried out with a
modification in which the task uncertainty is replaced by the average KL
divergence between all training tasks to each testing task. Due to the
extremely large number of training tasks, e.g. more than $10^{12}$ unique
5-way tasks can be generated from both the two datasets, the calculation of
the distance measure is infeasible. To make the training and testing tasks
manageable, we randomly generate $10,000$ tasks for training, and $100$ tasks
for testing. This results in $1,000,000$ distances, which can be calculated in
parallel with multiple computers. A testing task can be represented in the
correlation diagram through its prediction accuracy and the average KL
distance to training tasks, which is defined as:

$\overline{D}_{\mathrm{KL}}(\bm{\gamma}_{T+1})=\frac{1}{T}\sum_{i=1}^{T}D_{\mathrm{KL}}[q(\bm{\pi};\bm{\gamma}_{T+1})\|q(\bm{\pi};\bm{\gamma}_{i})].$

The correlation diagrams for both datasets are then plotted in Figures 5(c)
and 5(d). The results agree well with our hypothesis, in which the further a
testing task is from the training tasks, the worse the prediction accuracy.
This enables us to use the new representation produced by PTM to analyse task
similarity.

(a) Induction

(b) Transduction

(c) Final test results

Figure 6: Exponential weighted moving average (EWMA) of prediction accuracy
made by MAML following the lifelong learning for 100 random 5-way 1-shot tasks
sampled from mini-ImageNet testing set: 6(a) inductive setting, and 6(b)
transductive setting. The EWMA weight is set to 0.98 to smooth the noisy
signal. 6(c) Prediction accuracy made by models trained on different task
selection approaches on all 5-way 1-shot testing tasks generated from mini-
ImageNet. The error bars correspond to 95 percent confident interval.

###  4.2 Lifelong few-shot meta-learning

To further evaluate PTM, we conduct experiments following the lifelong
learning framework [29] with slight modification where the supervised tasks
are replaced by 5-way 1-shot learning episodes. More precisely, the setting
consists of a meta-learning model and a pool of $T_{\mathrm{pool}}$ tasks. At
each time step, a task selected from the pool is used to update the meta-
learning model, and discarded from the pool. A new task is then added to the
pool to maintain $T_{\mathrm{pool}}$ tasks available for learning. The
criterion for selecting a task to update the meta-learning model will depend
on the objective of interest. Two common objectives often observed in practice
are:

  * •

Induction: the selected training task is expected to encourage the meta-
learning model to be able to rapidly adapt to any future task,

  * •

Transduction: the selected training task is targeted toward one or many
specific testing tasks.

In the induction setting, the performance of the meta-learning model trained
on tasks selected by PTM is compared with three baselines: Task2Vec [1], the
worst-case approach [7] and random selection. For the PTM, the selection
criteria is based on the task entropy specified in Section 3, where the
training task with highest entropy is chosen for the learning. For Task2Vec,
tasks with large embedding norm are reported as difficult to learn. Hence, we
pick the one with the largest L1 norm produced by Task2Vec as the training
task. Originally, Task2Vec requires fine-tuning a pre-trained network (known
as probe network) on labelled data of a task. This fine-tuning step is,
however, infeasible for few-shot learning due to the insufficient number of
labelled data. We address this issue by training a MAML-based network to use
as a probe network. When given few-shot data of a training task, the MAML-
based probe network perform gradient update to adapt to that task. The task-
specific embedding can, therefore, be calculated using the adapted probe
network. We follow the Monte Carlo approach specified in the public code of
Task2Vec to calculate the corresponding task embedding. For the worst-case
approach, the training task that results in the highest loss for the current
meta-learning model is selected. Due to this nature, the worst-case approach
requires to evaluate all losses for each task in the pool at every time step,
leading to an extensive computation and might not scale well when the number
of tasks in the pool is large. For simplicity, we use MAML to train the meta-
learning model of interest for each selection strategy.

The transduction setting follows a similar setup as the induction case, but
the testing tasks, including the labelled and unlabelled data, are known
during training. For PTM, the average KL distances between all testing tasks
to each training task in the task pool are calculated, and the training task
with smallest average distance is selected. For Task2Vec, the proposed cosine
distance between normalised task embeddings is used to calculate the average
distance between all testing tasks to each training task [1]. Similar to PTM,
the training task with the smallest distance is prioritised for the learning.
For the worst-case approach, the entropy of the prediction $\hat{y}$ on
$C$-way testing tasks is used as the measure:

$S_{T+1}=-\sum_{c=1}^{C}\hat{y}_{c}\ln\hat{y}_{c},$

and the task that contributed to the highest entropy at prediction is chosen
[20]. The worst-case approach, therefore, requires $T_{\mathrm{pool}}$ trials
at every time step. In each trial, the current meta-model is adapted to each
training task in the pool, and then the average prediction entropy on all
testing tasks is calculated. This results in an extremely extensive
computation.

Four MAML-based meta-learning models are initialised identically and trained
on the tasks selected from a pool of $T_{\mathrm{pool}}=200$ tasks according
to the four criteria mentioned above. Figures 6(a) and 6(b) show the testing
results on 100 random mini-ImageNet tasks after every 500 time steps. Note
that the plotted results are smoothed by the exponential weighted moving
average with a weight of 0.98 to ease the visualisation. In general, PTM,
Task2Vec and worst-case can generalise better than random task selection. In
addition, the model trained with tasks chosen by PTM performs slightly better
than Task2Vec and the worst-case approach in both settings. This observation
might be explained based on the designated purpose of Task2Vec and the worst-
case approach. Task2Vec requires a sufficient number of labelled data to fine-
tune its probe network to calculate task embedding. Hence, it might not work
well in few-shot learning. For the worst-case, tasks are selected according to
a measure based on the current meta-model without taking task relatedness into
account. PTM, however, has a weakness in active selection since the approach
only focuses on task uncertainty or task similarity without considering the
current state of the meta-learning model. Nevertheless, PTM still provides a
good selection criterion comparing to Task2Vec and the worst-case approaches.
Note that although the active task selection is able to select the best task
within the pool, there might be the case where all remaining tasks in the pool
are uninformative, resulting in overfitting as observed in Figure 6(a).
However, for simplicity, no additional mechanism is integrated to decide
whether to learn from the selected task, or simply discarded from the pool. We
believe that adding L2 regularisation or applying early stopping based on a
validation set of tasks will help with this overfitting issue.

To further compare, we implement two additional meta-learning algorithms:
Prototypical Networks [31] and Amortised Bayesian meta-learning (ABML) [26]
and show results for the induction setting on all available testing 5-way
1-shot tasks of mini-ImageNet in Figure 6(c). Again, the prediction accuracy
made by the model trained on tasks selected by PTM outperforms other
baselines, especially the random one by a large margin.

##  5 Conclusion

We propose a generative approach based on variational auto-encoding and LDA
adopted in topic modelling to model tasks used in meta-learning. Under this
modelling approach, the dataset associated with a task can be expressed as a
mixture model of finite Gaussian distributions, where each task differs at the
mixture coefficients. An online VI method is presented to infer the parameters
of the Gaussian task-theme distributions. The obtained model allows us to
represent a task by its variational distribution of mixture coefficient in a
latent task-theme simplex, enabling the quantification of either the task
uncertainty or task similarity for active task selection.

## Acknowledgement

This work was supported by Australian Research Council through grants
DP180103232 and FT190100525. We also acknowledge the Phoenix HPC service at
the University of Adelaide that provided the super-computing resources for
this work.

## References

  * [1] Alessandro Achille, Michael Lam, Rahul Tewari, Avinash Ravichandran, Subhransu Maji, Charless C Fowlkes, Stefano Soatto and Pietro Perona  “TASK2VEC: Task embedding for meta-learning”  In _International Conference on Computer Vision_ , 2019, pp. 6430–6439
  * [2] Kelsey Allen, Evan Shelhamer, Hanul Shin and Joshua Tenenbaum  “Infinite mixture prototypes for few-shot learning”  In _International Conference on Machine Learning_ , 2019, pp. 232–241  PMLR
  * [3] Bart Bakker and Tom Heskes  “Task clustering and gating for Bayesian multitask learning”  In _Journal of Machine Learning Research_ 4.May, 2003, pp. 83–99
  * [4] Christopher M Bishop  “Pattern recognition and machine learning”  Springer, 2006
  * [5] David M Blei, Andrew Y Ng and Michael I Jordan  “Latent Dirichlet allocation”  In _Journal of Machine Learning Research_ 3.Jan, 2003, pp. 993–1022
  * [6] Kevin Canini, Lei Shi and Thomas Griffiths  “Online inference of topics with latent Dirichlet allocation”  In _Artificial Intelligence and Statistics_ , 2009, pp. 65–72
  * [7] Liam Collins, Aryan Mokhtari and Sanjay Shakkottai  “Task-Robust Model-Agnostic Meta-Learning”  In _Advances in Neural Information Processing Systems_ , 2020
  * [8] Rajarshi Das, Manzil Zaheer and Chris Dyer  “Gaussian LDA for topic models with word embeddings”  In _Annual Meeting of the Association for Computational Linguistics and the 7th International Joint Conference on Natural Language Processing (Volume 1: Long Papers)_ , 2015, pp. 795–804
  * [9] Guneet S Dhillon, Pratik Chaudhari, Avinash Ravichandran and Stefano Soatto  “A baseline for few-shot image classification”  In _International Conference on Learning Representations_ , 2020
  * [10] Chelsea Finn, Pieter Abbeel and Sergey Levine  “Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks”  In _International Conference on Machine Learning_ , 2017, pp. 1126–1135
  * [11] James Foulds, Levi Boyles, Christopher DuBois, Padhraic Smyth and Max Welling  “Stochastic collapsed variational Bayesian inference for latent Dirichlet allocation”  In _International Conference on Knowledge Discovery and Data Mining (ACM SIGKDD)_ , 2013, pp. 446–454
  * [12] Thomas L Griffiths and Mark Steyvers  “Finding scientific topics”  In _Proceedings of the National Academy of Sciences_ 101.suppl 1, 2004, pp. 5228–5235
  * [13] Matthew Hoffman, Francis R Bach and David M Blei  “Online learning for latent Dirichlet allocation”  In _Advances in Neural Information Processing Systems_ , 2010, pp. 856–864
  * [14] Laurent Jacob, Jean-philippe Vert and Francis R Bach  “Clustered multi-task learning: A convex formulation”  In _Advances in Neural Information Processing Systems_ , 2009, pp. 745–752
  * [15] Jean Kaddour and Steindór Sæmundsson  “Probabilistic active meta-learning”  In _Advances in Neural Information Processing Systems_ 33, 2020, pp. 20813–20822
  * [16] Diederik P Kingma and Max Welling  “Auto-encoding variational Bayes”  In _International Conference on Learning Representations_ , 2014
  * [17] Brenden M Lake, Ruslan Salakhutdinov and Joshua B Tenenbaum  “Human-level concept learning through probabilistic program induction”  In _Science_ 350.6266 American Association for the Advancement of Science, 2015, pp. 1332–1338
  * [18] Fei-Fei Li and Pietro Perona  “A Bayesian hierarchical model for learning natural scene categories”  In _International Conference on Computer Vision and Pattern Recognition_ 2, 2005, pp. 524–531
  * [19] Gabriel Loaiza-Ganem and John P Cunningham  “The continuous Bernoulli: fixing a pervasive error in variational autoencoders”  In _Advances in Neural Information Processing Systems_ , 2019
  * [20] David JC MacKay  “The evidence framework applied to classification networks”  In _Neural computation_ 4.5 MIT Press, 1992, pp. 720–736
  * [21] Thomas Minka  “Estimating a Dirichlet distribution”  Technical report, MIT, 2000
  * [22] Cuong Nguyen, Thanh-Toan Do and Gustavo Carneiro  “Similarity of Classification Tasks”  In _NeurIPS Meta-learning Workshop_ , 2020
  * [23] Cuong Nguyen, Thanh-Toan Do and Gustavo Carneiro  “Uncertainty in model-agnostic meta-learning using variational inference”  In _Winter Conference on Applications of Computer Vision_ , 2020, pp. 3090–3100
  * [24] Cuong V Nguyen, Tal Hassner, Cedric Archambeau and Matthias Seeger  “LEEP: A New Measure to Evaluate Transferability of Learned Representations”  In _International Conference on Machine Learning_ , 2020
  * [25] Jonathan K Pritchard, Matthew Stephens and Peter Donnelly  “Inference of population structure using multilocus genotype data”  In _Genetics_ 155.2 Genetics Soc America, 2000, pp. 945–959
  * [26] Sachin Ravi and Alex Beatson  “Amortized Bayesian Meta-Learning.”  In _International Conference on Learning Representations_ , 2019
  * [27] Sachin Ravi and Hugo Larochelle  “Optimization as a model for few-shot learning”  In _International Conference on Learning Representations_ , 2017
  * [28] Andrei A Rusu, Dushyant Rao, Jakub Sygnowski, Oriol Vinyals, Razvan Pascanu, Simon Osindero and Raia Hadsell  “Meta-learning with latent embedding optimization”  In _International Conference on Learning Representations_ , 2019
  * [29] Paul Ruvolo and Eric Eaton  “Active task selection for lifelong machine learning”  In _AAAI Conference on Artificial Intelligence_ 27, 2013
  * [30] Changjian Shui, Mahdieh Abbasi, Louis-Emile Robitaille, Boyu Wang and Christian Gagné  “A Principled Approach for Learning Task Similarity in Multitask Learning”  In _International Joint Conference on Artificial Intelligence_ , 2019, pp. 3446–3452
  * [31] Jake Snell, Kevin Swersky and Richard Zemel  “Prototypical networks for few-shot learning”  In _Advances in Neural Information Processing Systems_ , 2017, pp. 4077–4087
  * [32] Yee W Teh, David Newman and Max Welling  “A collapsed variational Bayesian inference algorithm for latent Dirichlet allocation”  In _Advances in neural information processing systems_ , 2007, pp. 1353–1360
  * [33] Sebastian Thrun and Joseph O’Sullivan  “Discovering structure in multiple learning tasks: The TC algorithm”  In _International Conference on Machine Learning_ 96, 1996, pp. 489–497
  * [34] Anh T Tran, Cuong V Nguyen and Tal Hassner  “Transferability and hardness of supervised classification tasks”  In _International Conference on Computer Vision_ , 2019, pp. 1395–1405
  * [35] Amos Tversky  “Features of similarity.”  In _Psychological review_ 84.4 American Psychological Association, 1977, pp. 327
  * [36] Oriol Vinyals, Charles Blundell, Tim Lillicrap and Daan Wierstra  “Matching networks for one shot learning”  In _Advances in Neural Information Processing Systems_ , 2016, pp. 3630–3638
  * [37] Ya Xue, Xuejun Liao, Lawrence Carin and Balaji Krishnapuram  “Multi-task learning for classification with Dirichlet process priors”  In _Journal of Machine Learning Research_ 8.Jan, 2007, pp. 35–63
  * [38] Jaesik Yoon, Taesup Kim, Ousmane Dia, Sungwoong Kim, Yoshua Bengio and Sungjin Ahn  “Bayesian Model-Agnostic Meta-Learning”  In _Advances in Neural Information Processing Systems_ , 2018, pp. 7343–7353
  * [39] Amir R Zamir, Alexander Sax, William Shen, Leonidas J Guibas, Jitendra Malik and Silvio Savarese  “Taskonomy: Disentangling task transfer learning”  In _Conference on Computer Vision and Pattern Recognition_ , 2018, pp. 3712–3722
  * [40] Yu Zhang and Dit-Yan Yeung  “A convex formulation for learning task relationships in multi-task learning”  In _Conference on Uncertainty in Artificial Intelligence_ , 2012

##  Appendix A Calculation of each term in the ELBO

As described in section 3, the variational distributions for
$\mathbf{u},\mathbf{z}$ and $\bm{\pi}$ are:

$\displaystyle q(\mathbf{u}_{in};\phi)$
$\displaystyle=\mathcal{N}\left(\mathbf{u}_{in};\mathbf{m}_{in},\left(\mathbf{s}_{in}\right)^{2}\mathbf{I}\right)$
(4) $\displaystyle q(\bm{\pi}_{i};\bm{\gamma}_{i})$
$\displaystyle=\mathrm{Dirichlet}\left(\bm{\pi}_{i};\bm{\gamma}_{i}\right)$
(10) $\displaystyle q(\mathbf{z}_{in};\mathbf{r}_{in})$
$\displaystyle=\mathrm{Categorical}\left(\mathbf{z}_{in};\mathbf{r}_{in}\right).$
(11)

###  A.1
$\mathbb{E}_{q(\mathbf{u}_{i};\bm{\mu}_{u_{i}},\bm{\Sigma}_{u_{i}})}\mathbb{E}_{q(\mathbf{z}_{i},\bm{\pi}_{i})}\left[\ln
p(\mathbf{u}_{i}|\mathbf{z}_{i},\bm{\mu},\bm{\Sigma})\right]$

$\displaystyle\mathbb{E}_{q(\mathbf{z}_{i},\bm{\pi}_{i})}\left[\ln
p(\mathbf{u}_{i}|\mathbf{z}_{i},\bm{\mu},\bm{\Sigma})\right]$
$\displaystyle=\sum_{n=1}^{N}\sum_{k=1}^{K}r_{ink}\ln
p(\mathbf{u}_{in}|\bm{\mu}_{k},\bm{\Sigma}_{k})$ (16)
$\displaystyle=\sum_{n=1}^{N}\sum_{k=1}^{K}r_{ink}\ln\mathcal{N}(\mathbf{u}_{in}|\bm{\mu}_{k},\bm{\Sigma}_{k}).$

Hence:

$\displaystyle\mathbb{E}_{q(\mathbf{u}_{i};\bm{\mu}_{u_{i}},\bm{\Sigma}_{u_{i}})}\mathbb{E}_{q(\mathbf{z}_{i},\bm{\pi}_{i})}\left[\ln
p(\mathbf{u}_{i}|\mathbf{z}_{i},\bm{\mu},\bm{\Sigma})\right]$
$\displaystyle=\sum_{n=1}^{N}\sum_{k=1}^{K}r_{ink}\underbrace{\mathbb{E}_{q(\mathbf{u}_{i};\bm{\mu}_{u_{i}},\bm{\Sigma}_{u_{i}})}\left[\ln\mathcal{N}(\mathbf{u}_{i}|\bm{\mu}_{k},\bm{\Sigma}_{k})\right]}_{\text{cross-
entropy between 2 Gaussians}}$ (17)
$\displaystyle=\sum_{n=1}^{N}\sum_{k=1}^{K}r_{ink}\left[-\frac{1}{2}\mathrm{tr}(\bm{\Sigma}_{k}^{-1}\bm{\Sigma}_{u_{in}})+\ln\mathcal{N}(\bm{\mu}_{u_{in}};\bm{\mu}_{k},\bm{\Sigma}_{k})\right].$

###  A.2
$\mathbb{E}_{q(\mathbf{u}_{i};\bm{\mu}_{u_{i}},\bm{\Sigma}_{u_{i}})}\mathbb{E}_{q(\mathbf{z}_{i},\bm{\pi}_{i})}\left[\ln
p(\mathbf{z}_{i}|\bm{\pi}_{i})\right]$

$\displaystyle\mathbb{E}_{q(\mathbf{u}_{i};\bm{\mu}_{u_{i}},\bm{\Sigma}_{u_{i}})}\mathbb{E}_{q(\mathbf{z}_{i},\bm{\pi}_{i})}\left[\ln
p(\mathbf{z}_{i}|\bm{\pi}_{i})\right]$
$\displaystyle=\sum_{n=1}^{N}\sum_{k=1}^{K}r_{ink}\int\mathrm{Dir}_{K}(\bm{\pi}_{i};\bm{\gamma}_{i})\ln\pi_{ik}\differential{\pi_{ik}}$
(18) $\displaystyle=\sum_{n=1}^{N}\sum_{k=1}^{K}r_{ink}\ln\tilde{\pi}_{ik},$

where:

$\boxed{\ln\tilde{\pi}_{ik}=\psi(\gamma_{ik})-\psi\left(\sum_{j=1}^{K}\gamma_{ij}\right).}$
(19)

###  A.3
$\mathbb{E}_{q(\mathbf{u}_{i};\bm{\mu}_{u_{i}},\bm{\Sigma}_{u_{i}})}\mathbb{E}_{q(\mathbf{z}_{i},\bm{\pi}_{i})}\left[\ln
p(\bm{\pi}_{i}|\bm{\alpha})\right]$

$\displaystyle\mathbb{E}_{q(\mathbf{u}_{i};\bm{\mu}_{u_{i}},\bm{\Sigma}_{u_{i}})}\mathbb{E}_{q(\mathbf{z}_{i},\bm{\pi}_{i})}\left[\ln
p(\bm{\pi}_{i}|\bm{\alpha})\right]$
$\displaystyle=\mathbb{E}_{\mathrm{Dir}(\bm{\pi}_{i};\bm{\gamma}_{i})}\left[\ln\Gamma\left(\sum_{j=1}^{K}\alpha_{j}\right)-\left[\sum_{k=1}^{K}\ln\Gamma(\alpha_{k})-(\alpha_{k}-1)\ln\pi_{ik}\right]\right]$
(20)
$\displaystyle=\left[\ln\Gamma\left(\sum_{j=1}^{K}\alpha_{j}\right)-\sum_{k=1}^{K}\ln\Gamma(\alpha_{k})\right]+\sum_{k=1}^{K}(\alpha_{k}-1)\ln\tilde{\pi}_{ik}.$

###  A.4
$\mathbb{E}_{q(\mathbf{u}_{i};\bm{\mu}_{u_{i}},\bm{\Sigma}_{u_{i}})}\mathbb{E}_{q(\mathbf{z}_{i},\bm{\pi}_{i})}\left[\ln
q(\mathbf{z}_{i}|\mathbf{r}_{i})\right]$

$\mathbb{E}_{q(\mathbf{u}_{i};\bm{\mu}_{u_{i}},\bm{\Sigma}_{u_{i}})}\mathbb{E}_{q(\mathbf{z}_{i},\bm{\pi}_{i})}\left[\ln
q(\mathbf{z}_{i}|\mathbf{r}_{i})\right]=\sum_{n=1}^{N}\sum_{k=1}^{K}r_{ink}\ln
r_{ink}.$ (21)

###  A.5
$\mathbb{E}_{q(\mathbf{u}_{i};\bm{\mu}_{u_{i}},\bm{\Sigma}_{u_{i}})}\mathbb{E}_{q(\mathbf{z}_{i},\bm{\pi}_{i})}\left[\ln
q(\bm{\pi}_{i}|\bm{\gamma}_{i})\right]$

$\mathbb{E}_{q(\mathbf{u}_{i};\bm{\mu}_{u_{i}},\bm{\Sigma}_{u_{i}})}\mathbb{E}_{q(\mathbf{z}_{i},\bm{\pi}_{i})}\left[\ln
q(\bm{\pi}_{i}|\bm{\gamma}_{i})\right]=\ln\Gamma\left(\sum_{j=1}^{K}\gamma_{ij}\right)-\sum_{k=1}^{K}\left[\ln\Gamma(\gamma_{ik})-(\gamma_{ik}-1)\ln\tilde{\pi}_{ik}\right].$
(22)

##  Appendix B Maximisation of the ELBO

Since the ELBO can be evaluated as shown in Appendix A, we can maximise the
ELBO w.r.t. task-specific variational parameters by taking derivative, setting
it to zero and solving for the parameters of interest.

###  B.1 Variational categorical distribution

Note that:

$\sum_{k=1}^{K}r_{ink}=1.$ (23)

The derivative of $\mathsf{L}_{i}$ with respect to $r_{ink}$ can be expressed
as:

$\displaystyle\frac{\partial\mathsf{L}}{\partial
r_{ink}}=-\frac{1}{2}\mathrm{tr}(\bm{\Sigma}_{k}^{-1}\bm{\Sigma}_{u_{in}})+\ln\mathcal{N}(\bm{\mu}_{u_{in}};\bm{\mu}_{k},\bm{\Sigma}_{k})+\ln\tilde{\pi}_{ik}-\ln
r_{ink}-1+\lambda,$ (24)

where: $\lambda$ is the Lagrange multiplier and $\ln\tilde{\pi}_{ik}$ is
defined in Eq. (19). Setting the derivative to zero and solving for $r_{ink}$
give:

$\boxed{r_{ink}\propto\exp\left[-\frac{1}{2}\mathrm{tr}(\bm{\Sigma}_{k}^{-1}\bm{\Sigma}_{u_{in}})+\ln\mathcal{N}(\bm{\mu}_{u_{in}};\bm{\mu}_{k},\bm{\Sigma}_{k})+\ln\tilde{\pi}_{ik}\right].}$
(25)

###  B.2 Variational Dirichlet distribution

The lower-bound related to $\gamma_{ik}$ can be written as:

$\displaystyle\mathsf{L}$
$\displaystyle=\sum_{k=1}^{K}\sum_{n=1}^{N}r_{ink}\ln\tilde{\pi}_{ik}+\sum_{k=1}^{K}(\alpha_{k}-1)\ln\tilde{\pi}_{ik}-\ln\Gamma\left(\sum_{j=1}^{K}\gamma_{ij}\right)+\sum_{k=1}^{K}\left[\ln\Gamma(\gamma_{ik})-(\gamma_{ik}-1)\ln\tilde{\pi}_{ik}\right]$
(26)
$\displaystyle=-\ln\Gamma\left(\sum_{j=1}^{K}\gamma_{ij}\right)+\sum_{k=1}^{K}\ln\tilde{\pi}_{ik}\left(\alpha_{k}-\gamma_{ik}+\sum_{n=1}^{N}r_{ink}\right)+\ln\Gamma(\gamma_{ik})$
$\displaystyle=-\ln\Gamma\left(\sum_{j=1}^{K}\gamma_{ij}\right)+\sum_{k=1}^{K}\left[\psi(\gamma_{ik})-\psi\left(\sum_{j=1}^{K}\gamma_{ij}\right)\right]\left(\alpha_{k}-\gamma_{ik}+\sum_{n=1}^{N}r_{ink}\right)+\ln\Gamma(\gamma_{ik}).$

Hence, the lower-bound related to $\gamma_{ik}$ is:

$\displaystyle\mathsf{L}[\gamma_{ik}]$
$\displaystyle=-\ln\Gamma\left(\sum_{j=1}^{K}\gamma_{ij}\right)+\psi(\gamma_{ik})\left(\alpha_{k}-\gamma_{ik}+\sum_{n=1}^{N}r_{ink}\right)$
(27)
$\displaystyle\qquad-\psi\left(\sum_{j=1}^{K}\gamma_{ij}\right)\left(\sum_{j=1}^{K}\alpha_{j}-\gamma_{ij}+\sum_{n=1}^{N}r_{inj}\right)+\ln\Gamma(\gamma_{ik})$

Taking derivative w.r.t. $\gamma_{ik}$ gives:

$\displaystyle\frac{\partial\mathsf{L}}{\partial\gamma_{ik}}$
$\displaystyle=-\psi\left(\sum_{j=1}^{K}\gamma_{ij}\right)+\Psi(\gamma_{ik})\left(\alpha_{k}-\gamma_{ik}+\sum_{n=1}^{N}r_{ink}\right)-\psi(\gamma_{ik})$
(28)
$\displaystyle\quad-\Psi\left(\sum_{j=1}^{K}\right)\left(\sum_{j=1}^{K}\alpha_{j}-\gamma_{ij}+\sum_{n=1}^{N}r_{inj}\right)+\psi\left(\sum_{j=1}^{K}\gamma_{ij}\right)+\psi(\gamma_{ik})$
$\displaystyle=\Psi(\gamma_{ik})\left(\alpha_{k}-\gamma_{ik}+\sum_{n=1}^{N}r_{ink}\right)-\Psi\left(\sum_{j=1}^{K}\right)\sum_{j=1}^{K}\alpha_{j}-\gamma_{ij}+\sum_{n=1}^{N}r_{inj},$

where $\Psi(.)$ is the trigamma function.

Setting the derivative to zero yields a maximum at:

$\boxed{\gamma_{ik}=\alpha_{k}+N_{ik},}$ (29)

where:

$N_{ik}=\sum_{n=1}^{N}r_{ink}.$ (30)

###  B.3 Maximum likelihood for the task-theme $\bm{\mu}_{k}$ and
$\bm{\Sigma}_{k}$

The terms in the objective function relating to $\mu_{k}$ can be written as:

$\displaystyle\mathsf{L}[\mu_{k}]$
$\displaystyle=\sum_{i=1}^{T}\sum_{n=1}^{N}r_{ink}\ln\mathcal{N}\left(\bm{\mu}_{u_{in}};\bm{\mu}_{k},\bm{\Sigma}_{k}\right)$
(31)
$\displaystyle=-\frac{1}{2}\sum_{i=1}^{T}\sum_{n=1}^{N}r_{ink}\left(\bm{\mu}_{u_{in}}-\bm{\mu}_{k}\right)^{\top}\bm{\Sigma}_{k}^{-1}\left(\bm{\mu}_{u_{in}}-\bm{\mu}_{k}\right)$

Taking derivative w.r.t. $\bm{\mu}_{k}$ gives:

$\partialderivative{\mathsf{L}}{\bm{\mu}_{k}}=\sum_{i=1}^{T}\sum_{n=1}^{N}r_{ink}\bm{\Sigma}_{k}^{-1}(\bm{\mu}_{u_{in}}-\bm{\mu}_{k}).$
(32)

Setting the derivative to zero yields a maximum at:

$\boxed{\bm{\mu}_{k}=\frac{\sum_{i=1}^{T}\sum_{n=1}^{N}r_{ink}\bm{\mu}_{u_{in}}}{\sum_{i=1}^{T}N_{ik}}.}$
(33)

The terms in the objective function relating to $\bm{\Sigma}_{k}$ is given as:

$\displaystyle\mathsf{L}$
$\displaystyle=\sum_{i=1}^{T}\sum_{n=1}^{N}r_{ink}\left[-\frac{1}{2}\mathrm{tr}(\bm{\Sigma}_{k}^{-1}\bm{\Sigma}_{u_{in}})+\ln\mathcal{N}(\bm{\mu}_{u_{in}};\bm{\mu}_{k},\bm{\Sigma}_{k})\right]$
(34)
$\displaystyle=-\frac{1}{2}\sum_{i=1}^{T}\sum_{n=1}^{N}r_{ink}\left[\mathrm{tr}(\bm{\Sigma}_{k}^{-1}\bm{\Sigma}_{u_{in}})+d\ln(2\pi)+\ln|\bm{\Sigma}_{k}|+\left(\bm{\mu}_{u_{in}}-\bm{\mu}_{k}\right)^{\top}\bm{\Sigma}_{k}^{-1}\left(\bm{\mu}_{u_{in}}-\bm{\mu}_{k}\right)\right].$

Taking derivative w.r.t. $\bm{\Sigma}_{k}$ gives:

$\displaystyle\partialderivative{\mathsf{L}}{\bm{\Sigma}_{k}}$
$\displaystyle=-\frac{1}{2}\sum_{i=1}^{T}\sum_{n=1}^{N}r_{ink}\left[-\bm{\Sigma}_{k}^{-1}\bm{\Sigma}_{u_{in}}\bm{\Sigma}_{k}^{-1}+\bm{\Sigma}_{k}^{-1}-\bm{\Sigma}_{k}^{-1}\left(\bm{\mu}_{u_{in}}-\bm{\mu}_{k}\right)\left(\bm{\mu}_{u_{in}}-\bm{\mu}_{k}\right)^{\top}\bm{\Sigma}_{k}^{-1}\right]$
(35)
$\displaystyle=\frac{1}{2}\sum_{i=1}^{T}\sum_{n=1}^{N}r_{ink}\left[\bm{\Sigma}_{k}^{-1}\bm{\Sigma}_{u_{in}}-\mathbf{I}+\bm{\Sigma}_{k}^{-1}\left(\bm{\mu}_{u_{in}}-\bm{\mu}_{k}\right)\left(\bm{\mu}_{u_{in}}-\bm{\mu}_{k}\right)^{\top}\right]\bm{\Sigma}_{k}^{-1}.$

Setting the derivative to zero gives:

$\boxed{\begin{aligned}
\bm{\Sigma}_{k}&=\frac{1}{\sum_{i=1}^{T}N_{ik}}\sum_{i=1}^{T}\sum_{n=1}^{N}r_{ink}\left[\bm{\Sigma}_{u_{in}}+\left(\bm{\mu}_{u_{in}}-\bm{\mu}_{k}\right)\left(\bm{\mu}_{u_{in}}-\bm{\mu}_{k}\right)^{\top}\right].\end{aligned}}$
(36)

###  B.4 Maximum likelihood for $\bm{\alpha}$

The lower-bound with terms relating to $\alpha_{k}$ can be expressed as:

$\mathsf{L}=T\left[\ln\Gamma\left(\sum_{j=1}^{K}\alpha_{j}\right)-\sum_{k=1}^{K}\ln\Gamma(\alpha_{k})\right]+\sum_{i=1}^{T}\sum_{k=1}^{K}(\alpha_{k}-1)\left[\psi(\gamma_{ik})-\psi\left(\sum_{j=1}^{K}\gamma_{ij}\right)\right].$
(37)

Taking derivative w.r.t. $\alpha_{k}$ gives:

$\displaystyle g_{k}=\partialderivative{\mathsf{L}}{\alpha_{k}}$
$\displaystyle=T\left[\psi\left(\sum_{j=1}^{K}\alpha_{j}\right)-\psi(\alpha_{k})\right]+\sum_{i=1}^{T}\left[\psi(\gamma_{ik})-\psi\left(\sum_{j=1}^{K}\gamma_{ij}\right)\right].$
(38)

The second derivative is, therefore, obtained as:

$\displaystyle\frac{\partial^{2}\mathsf{L}}{\partial\alpha_{k}\partial\alpha_{k^{\prime}}}$
$\displaystyle=T\left[\Psi\left(\sum_{j=1}^{K}\alpha_{j}\right)-\delta(k-k^{\prime})\Psi(\alpha_{k})\right].$
(39)

The Hessian can be written in matrix form [21] as:

$\displaystyle\mathbf{H}$ $\displaystyle=\mathbf{Q}+\mathbf{11}^{T}a$ (40)
$\displaystyle q_{kk^{\prime}}$
$\displaystyle=-T\delta(k-k^{\prime})\Psi(\alpha_{k})$ (41) $\displaystyle a$
$\displaystyle=T\Psi\left(\sum_{j=1}^{K}\alpha_{j}\right).$ (42)

One Newton step is therefore:

$\displaystyle\bm{\alpha}$
$\displaystyle\leftarrow\bm{\alpha}-\mathbf{H}^{-1}\mathbf{g}$ (43)
$\displaystyle(\mathbf{H}^{-1}\mathbf{g})_{k}$
$\displaystyle=\frac{g_{k}-b}{q_{kk}},$ (44)

where:

$b=\frac{\sum_{j=1}^{K}g_{j}/q_{jj}}{1/a+\sum_{j=1}^{K}1/q_{jj}}.$ (45)