# Omni-Training: Bridging Pre-Training and Meta-Training for Few-Shot Learning

Yang Shu, Zhangjie Cao, Jinghan Gao, Jianmin Wang, Philip S. Yu, , Mingsheng
Long  The authors are with the School of Software, BNRist, Tsinghua
University, Beijing 100084, China. Corresponding author: Mingsheng Long,
<EMAIL_ADDRESS>
###### Abstract

Few-shot learning aims to fast adapt a deep model from a few examples. While
pre-training and meta-training can create deep models powerful for few-shot
generalization, we find that pre-training and meta-training focuses
respectively on cross-domain transferability and cross-task transferability,
which restricts their data efficiency in the entangled settings of domain
shift and task shift. We thus propose the Omni-Training framework to
seamlessly bridge pre-training and meta-training for data-efficient few-shot
learning. Our first contribution is a tri-flow Omni-Net architecture. Besides
the joint representation flow, Omni-Net introduces two parallel flows for pre-
training and meta-training, responsible for improving domain transferability
and task transferability respectively. Omni-Net further coordinates the
parallel flows by routing their representations via the joint-flow, enabling
knowledge transfer across flows. Our second contribution is the Omni-Loss,
which introduces a self-distillation strategy separately on the pre-training
and meta-training objectives for boosting knowledge transfer throughout
different training stages. Omni-Training is a general framework to accommodate
many existing algorithms. Evaluations justify that our single framework
consistently and clearly outperforms the individual state-of-the-art methods
on both cross-task and cross-domain settings in a variety of classification,
regression and reinforcement learning problems.

###### Index Terms:

Few-shot learning, data efficiency, transferability, meta-learning, pre-
training

##  1 Introduction

Deep learning [1] has achieved the state-of-the-art performance in various
machine learning tasks [2, 3, 4, 5]. However, most deep learning methods, in
particular the foundation models [6], are “data hungry”, in that the success
of these methods highly relies on large amounts of labeled data. This clearly
limits the application of deep learning to widespread domains or tasks,
especially those with sparse data and insufficient annotations, such as
personalized healthcare [7]. In order to promote the grounding of deep
learning models, few-shot learning, which aims to fast learn various complex
tasks from a few labeled data, has attracted enormous attention recently [8,
9, 10].

Human beings are gifted with the ability to quickly learn new tasks by making
use of previous experience and knowledge. In analogy to this, deep learning
models can reuse the representations learned previously to help efficiently
solve widespread downstream tasks. Recent advances have revealed that a
properly trained model endows an important property: _transferability_ , and
higher transferability indicates better generalizability to new scenarios. In
general situations as illustrated by Figure 1, complex relationships between
the pretext dataset and the new task hinder the downstream learning and pose
challenges to the transferability of learned representations. The two main
challenges come from the different distributions across domains, _i.e._ domain
shift and different semantics across tasks, _i.e._ task shift. For example, in
image classification, different domains may have different visual factors such
as different styles, viewpoints and lighting, while different tasks may have
different categories. In most cases, the two challenges entangle with each
other, making few-shot learning a very hard problem. Thus, a versatile
algorithm should bridge these two gaps and learn representations with both
_domain transferability_ and _task transferability_.

Figure 1: Illustration of the two challenges of few-shot learning. Due to the
_domain shift_ and _task shift_ between the training dataset and the test
dataset, it is hard for the trained model $\mathcal{M}_{\texttt{train}}$ to
transfer to the test set and boost its data-efficiency. An ideal training
method should learn representations with both _domain transferability_ and
_task transferability_ and adapt $\mathcal{M}_{\texttt{train}}$ to downstream
model $\mathcal{M}_{\texttt{test}}$ in a data-efficient way.

Two mainstream representation learning paradigms for few-shot learning are
_pre-training_ and _meta-training_. In pre-training, we train a high-capacity
model for a pretext task on large-scale datasets [4, 11] and fine-tune the
model on the target task [12]. In meta-training, we train the model from
diverse tasks and fast adapt the model to new tasks [13, 9, 14]. As evidenced
by recent studies, neither paradigm can dominate in the widespread few-shot
learning scenarios [15, 16, 17], because it requires knowledge that
generalizes across both domains and tasks. Pre-training representations can
transfer to widespread domains, since the pretext task is designed to be
general across domains. However, only pre-training on a single pretext task
makes it hard to fast adapt to many new tasks. In contrast, the diverse tasks
equip meta-training with the ability to fast adapt across many tasks with
extremely sparse data, but the meta-training tasks are usually domain-specific
and thus the learned representations cannot generalize well across domains.

In line with the understanding of pre-training and meta-training, we further
study both paradigms with regard to the two transferability properties and
reach a similar conclusion: pre-training methods are apt at the domain
transferability while meta-training methods at the task transferability. We
then take a step forward to exploit the collaboration between pre-training and
meta-training and draw an important finding that neither a simple ensemble nor
a tight combination can achieve both kinds of transferability. This finding
motivates us to design a new Omni-Training framework to bridge both sides for
few-shot learning.

Omni-Training seamlessly bridges pre-training and meta-training to learn deep
representations with both domain transferability and task transferability. The
first part is Omni-Net, a tri-flow architecture. Besides a joint-flow for
shared representation learning, Omni-Net introduces two new parallel flows for
pre-training and meta-training to yield representations of domain
transferability and task transferability respectively. It further coordinates
the parallel flows by routing their representations via the joint-flow, making
each gain the other kind of transferability. The second part is Omni-Loss,
which works in cooperation with the architecture for learning transferable
representations. A self-distillation strategy is imposed to both the pre-
training and meta-training objectives, forcing the parallel flows to learn
more transferable representations. Omni-Training is a general framework that
can accommodate many existing pre-training and meta-training algorithms.
Thorough evaluations on cross-task and cross-domain datasets in
classification, regression and reinforcement learning problems show that Omni-
Training consistently and clearly outperforms the individual state-of-the-art
deep learning methods.

##  2 Related Work

Few-shot learning aims to make full use of every sample and address new tasks
with a few labeled data [8, 18, 19]. In this paper, we focus on representation
learning algorithms towards it, which aim to learn transferable
representations from pretext data to reduce the data requirement of learning
new tasks. We restrict our review to two mainstream categories of
representation learning algorithms for few-shot learning that achieve state-
of-the-art performance: pre-training and meta-training.

###  2.1 Pre-Training

One line of few-shot learning methods is to learn deep representations by pre-
training deep networks with a pretext task on the training datasets. With the
prevalence of large-scale labeled datasets and the advanced computational
infrastructure, deep networks with extremely big model capacity are trained
for various applications such as computer vision [20, 21, 11] and natural
language processing [4, 22]. With such deep models, recent works re-take the
pre-training and fine-tuning paradigm and demonstrate that fine-tuning high-
capacity deep models pre-trained on large datasets achieves state-of-the-art
performance in various applications with only a few labeled data [23, 24, 25].
Pre-training is also adopted in reinforcement learning to enable learning the
policy for new environments with less interaction steps [26, 27, 28]. More
advanced pre-training strategies also boost few-shot learning performance,
such as training an ensemble of models [29] and training with knowledge
distillation [30].

There are methods towards the stage of fine-tuning on the new task. For
example, some works reuse the representations to predict parameters of new
categories [31, 32]. Some works regularize the model of the new task from the
aspects of parameters or representations to fully extract the knowledge of the
pre-trained models [33, 34]. Recent research also proposed to explore
relationships between the training and test datasets and mitigate negative
transfer [35, 36]. Cao et al. [24] proposed an ease-in-ease-out fine-tuning
method to enable transfer reinforcement learning across homotopy classes.
These methods focus on a different perspective and are in parallel with this
paper.

Pre-training approaches are simple and effective to improve data efficiency in
new scenarios, which show higher domain transferability and outperform
sophisticated meta-training methods in the cross-domain setting [15, 16, 17].
However, as the training stage only involves one pretext task, these methods
cannot quickly handle the rapid changes of semantics in new tasks [9].

###  2.2 Meta-Training

Meta-training addresses few-shot learning by learning representations
generalizable across many training tasks, which can be naturally adapted to
new tasks [37, 38, 39]. It has been widely used in a variety of applications.

Few-shot learning [8] is widely studied in the field of classification,
especially image recognition, where a typical form is to learn from a few
annotated data, _i.e._ the “N-way-K-shot” few-shot classification problems
[18, 40]. Metric-based meta-learning methods are tailored for these problems,
which learn an embedding space to form decision boundaries according to the
distances between samples [41, 13, 14, 42, 43]. Recently, embedding functions
are improved by stronger inductive bias such as graph networks [44], fine-
grained attention maps [45], task-adaptive projections [46, 47] and set-to-set
functions [48].

Some other meta-learning methods deal with various applications. Early works
build meta-learners to learn how to update the model parameters and generalize
the updating rules to new tasks [49, 50], which have been recently applied in
deep learning to enable fast adaptation of deep networks [51, 52, 53]. Such
learning to learn paradigm is also demonstrated to work for regression [51,
53] and reinforcement learning [54, 55]. Several works equip networks with
external or internal memory so that meta-knowledge can be effectively stored
and queried for data-efficient adaptation to new tasks [56, 57, 58, 59]. The
memory-augmented models are also applied to reinforcement learning to improve
data-efficiency [60, 61, 58]. These methods introduce additional parameters
and storage costs or require a particular architecture of the learner for
meta-learning.

Model agnostic meta-learning introduces the gradient-based idea, which trains
a good initialization of the deep network as the meta-knowledge such that a
small number of gradient steps and interactions in the new environment can
induce high generalization performance [9]. The idea is later improved by new
architectures [62, 63]. Such gradient-based meta-training methods show strong
performance in real robotics applications such as imitation learning [64, 65],
locomotion [66], visual navigation [67], and robot manipulation [68]. They can
also be extended to other applications such as regression and image
classification by changing the architecture and training objective [9, 69, 70,
71].

Though meta-training empowers the deep representations with the ability to
generalize across new tasks, a recent empirical study has revealed that meta-
trained representations cannot generalize across domains with distribution
shift [15]. Tseng et al. [72] use feature-wise transformation layers to
simulate various image feature distributions extracted from the training tasks
in different domains. However, the domain transferability is still limited
especially in domains with large distribution shift [16]. Our method acquires
the missing piece of domain transferability from _pre-training_ , which does
not require multiple pretext domains but achieves better cross-domain
generalization ability.

Meta-training and pre-training are apt at task transferability and domain
transferability respectively, and neither can dominate the other. A natural
idea is to integrate two types of approaches to achieve both. Sun et al. [73]
simply chain the process of pre-training and meta-training, but such a simple
combination still lacks both kinds of transferability. In contrast, our Omni-
Training framework seeks to flexibly bridge pre-training and meta-training to
empower both kinds of transferability.

##  3 Background and Analysis

We first introduce few-shot learning and its two key prerequisites: _domain
transferability_ and _task transferability_. Then we delve into two mainstream
methods, _pre-training_ and _meta-training_ , each of which learns a
representation of a specific kind of transferability and enables
generalization to either new domains or new tasks.

###  3.1 Few-Shot Learning

At the training phase, the goal is to learn a feature representation $F$ on
the training set $\mathcal{D}_{\texttt{train}}$ of sufficient labeled
examples, which enables fast solving new tasks from a few examples. At the
testing phase, the learned representation $F$ is evaluated on new tasks,
either within domain or across domains. Each task comes with a test set
$\mathcal{D}_{\texttt{test}}=\\{\mathcal{S}_{\texttt{test}},\mathcal{Q}_{\texttt{test}}\\}$
partitioned into a support set $\mathcal{S}_{\texttt{test}}$ with a few
labeled examples and a query set $\mathcal{Q}_{\texttt{test}}$ with many
unlabeled examples to predict. The learned representation $F$ should adapt
fast to each new task through the support set and then yield accurate
predictions on the query set.

The key to enable few-shot learning in downstream tasks is the transferability
of the representations. Given input $\mathbf{x}\in\mathcal{X}$ and output
$\mathbf{y}\in\mathcal{Y}$, denote the joint distribution as
$\mathcal{P}(\mathbf{x},\mathbf{y})$ and the learning task as
$\mathcal{T}:\mathbf{x}\mapsto\mathbf{y}$. The _domain transferability_
measures the generalizability under train-test distribution shift,
$\mathcal{P}_{\texttt{train}}\neq\mathcal{P}_{\texttt{test}}$, and the _task
transferability_ measures the generalizability under train-test task shift,
$\mathcal{T}_{\texttt{train}}\neq\mathcal{T}_{\texttt{test}}$. In general
situations of few-shot learning, complex relationships between the training
dataset and the new tasks entangle distribution shift and task shift. So we
should learn representations with both domain transferability and task
transferability to enable data-efficient few-shot learning.

(a) Task Transferability

(b) Domain Transferability

(c) Combination Strategies

Figure 2: Analysis: (a) Task transferability of pre-training and meta-training
in the _cross-task_ setting; (b) Domain transferability of pre-training and
meta-training in the _cross-domain_ setting; (c) Accuracy of pre-training,
meta-training and two combination strategies, Ensemble and Joint-Training.

###  3.2 Training Methods for Few-Shot Learning

Pre-Training. In pre-training approaches, deep representations are often
learned by supervised learning on a large-scale training dataset
$\mathcal{D}_{\texttt{train}}$, which facilitate data-efficient or few-shot
learning for a variety of downstream tasks. We use an abstract model composed
of a feature extractor $F$ to generate the representation and a task-specific
head $H$ to predict the output, which is applicable to various tasks. During
the training stage, the training set $\mathcal{D}_{\texttt{train}}$ is viewed
as samples from a joint distribution of inputs and labels:
$\mathcal{P}(\mathbf{x},\mathbf{y})$. Representation learning is conducted by
optimizing $H$ and $F$ over the sampled _mini-batches_ from the training
distribution with the loss $\ell_{\texttt{pre}}$ tailored to the specific task
or algorithm:

$\mathop{\min}_{H,F}{\mathbb{E}}_{(\mathbf{x},\mathbf{y})\sim\mathcal{P}(\mathbf{x},\mathbf{y})}\
\ell_{\texttt{pre}}\big{(}\mathbf{y},H\circ F\left(\mathbf{x}\right)\big{)}.$
(1)

During test, we transfer the pre-trained models on the new task
$\mathcal{D}_{\texttt{test}}=\\{\mathcal{S}_{\texttt{test}},\mathcal{Q}_{\texttt{test}}\\}$.
The feature extractor $F$ is fine-tuned and a task-specific head
$H_{\texttt{new}}$ for the new task is trained with the labeled data in
support set $\mathcal{S}_{\texttt{test}}$ and applied in query set
$\mathcal{Q}_{\texttt{test}}$.

Meta-Training. In meta-training, the representations are learned to perform
well across a set of tasks sampled from a task distribution constructed from
the training set. Specifically, the training set
$\mathcal{D}_{\texttt{train}}$ is viewed as a distribution of tasks
$\mathcal{T}$. Each task mimics the testing situation, which contains a
support set $\mathcal{S}$ with only a few labeled samples and a query set
$\mathcal{Q}$ needing predictions. The meta-learner is optimized over
_episodes_ of tasks sampled from $\mathcal{T}$. The model $F$ and $H$ are
learned to efficiently solve each of the tasks conditioned on the support set
$\mathcal{S}$ with only a few samples, and updated by the performance
evaluated on the query set $\mathcal{Q}$:

$\mathop{\min}_{H,F}\mathbb{E}_{(\mathcal{S},\mathcal{Q})\sim\mathcal{P}(\mathcal{T})}\mathbb{E}_{(\mathbf{x},\mathbf{y})\in\mathcal{Q}}\
\ell_{\texttt{meta}}\big{(}\mathbf{y},H\circ
F(\mathbf{x}|\mathcal{S})\big{)},$ (2)

where $\ell_{\texttt{meta}}$ is the loss of specific meta-training algorithms
defined on each episode, _e.g._ , the meta-objective in [9]. In test time, the
models are fast adapted to the new task with its support set
$\mathcal{S}_{\texttt{test}}$ in a similar way as the training phase, and the
adapted models can be used for predictions on the query set
$\mathcal{Q}_{\texttt{test}}$.

###  3.3 Transferability Assessment

We empirically compare pre-training and meta-training in terms of task
transferability and domain transferability. We evaluate two typical methods,
Baseline [15] as the pre-training method and ProtoNet [14] as the meta-
training method. We first use two benchmarks mini-ImageNet and CUB and follow
the protocol in [15]. Note that we use test tasks from the same dataset to
relieve the influence of distribution shift and mainly focus on _task shift_.
As shown in Figure 2a, the pre-training and meta-training methods perform
comparably on mini-ImageNet-5 (5 examples per class). However, in the more
extreme situation with only 1 example per class, meta-training outperforms
pre-training, where the boost becomes larger on CUB: a fine-grained dataset
with smaller distribution shifts between tasks. The result indicates higher
task transferability of meta-training. Next, we explore the influence of
_distribution shift_ across domains. We train the model on the mini-ImageNet
dataset, but evaluate it on different domains including CUB, Cars, Places and
Plantae. As shown in Figure 2b, pre-training and meta-training have similar
in-domain performance, but pre-training consistently outperforms meta-training
in four cross-domain situations. This result indicates higher domain
transferability of pre-training.

Our key finding is that pre-training introduces higher domain transferability
while meta-training introduces higher task transferability. This explains the
phenomenon that both methods may fail in some few-shot learning scenarios [9,
15, 74, 16]. In general situations, the new tasks hold complex relationships
with the training set, presenting both challenges of distribution shift and
task shift, which entangle with each other. For example, in the in-domain
experiment, there could still be domain shift caused by different categories;
In the cross-domain experiment, while domain shift is the main challenge, task
transferability is still required to adapt across different classes. Overall,
we need to learn representations with both domain transferability and task
transferability to fully enable few-shot learning.

We study two simple ways to combine pre-training and meta-training. One is to
separately train two models with two methods, and use their ensemble for
prediction, denoted as Ensemble. The other is to jointly train the model with
both training objectives, denoted as Joint-Training. We evaluate them on three
situations of mini-ImageNet, CUB, and transferring mini-ImageNet to CUB. As
shown in Figure 2c, both combination strategies promote the performance in
some cases, but the improvement is minor and inconsistent. The gain of
Ensemble indicates that pre-training and meta-training representations endow
complementary knowledge. However, this simple ensemble lacks the knowledge
coordination between pre-training and meta-training. The improvement of Joint-
Training shows the importance to extract shared knowledge between the two
training paradigms, but this tight combination sacrifices the specific
transferability held by each approach. Such a _transferability dilemma_
motivates the proposed Omni-Training framework, which seeks to flexibly
acquire both domain transferability and task transferability for better few-
shot learning.

##  4 Omni-Training Framework

In this paper, we are interested in learning representations with both domain
transferability and task transferability by incorporating and bridging pre-
training and meta-training in a unified Omni-Training framework. As discussed
in Section 3.3, this goal is non-trivial to realize with simple combinations
of these two training paradigms. Beyond the tight combination of joint-
training, we have two more key insights in designing the framework. Our first
key insight is that the domain transferability of pre-training and the task
transferability of meta-training should be preserved. Furthermore, there
should be knowledge communication between the two types of training to enable
them to complement each other. Our second key insight is that this non-trivial
unification should be realized with the design in both network architectures
and training algorithms. These insights are embedded into the Omni-Training
framework via an Omni-Net architecture guided by an Omni-Loss.

Figure 3: The Omni-Training framework consists of three data flows: joint-flow
(green), pre-flow (blue), and meta-flow (red). The Omni-Net consists of a
backbone $F$ and an Omni-Head $H$, where $F$ is formed by stacking Omni-Layers
and $H$ is formed of three heads $H_{\texttt{joint}}$, $H_{\texttt{pre}}$ and
$H_{\texttt{meta}}$. Each Omni-Layer has a main chunk layer
$f_{\texttt{joint}}$ and two lightweight branch layers $f_{\texttt{pre}}$ and
$f_{\texttt{meta}}$, followed by activation functions $a_{\texttt{joint}}$,
$a_{\texttt{pre}}$, $a_{\texttt{meta}}$. The Omni-Loss consists of three
losses respectively for joint-training $\mathcal{L}_{\texttt{joint}}$, pre-
training $\mathcal{J}_{\texttt{pre}}$, and meta-training
$\mathcal{J}_{\texttt{meta}}$, computed on the corresponding head. We also
propose a self-distillation strategy for training the pre-flow and meta-flow,
which transfers knowledge throughout the training process.

###  4.1 Omni-Net

Omni-Net is a tri-flow architecture that is constructed by stacking Omni-
Layers for representation learning and Omni-Heads for output prediction, as
shown in Figure 3.

Omni-Layer. We aim to simultaneously preserve the domain transferability of
pre-training and the task transferability of meta-training, and promote
knowledge communication between them. Thus, as shown in Figure 3, we design an
Omni-Layer consisting of a main chunk layer $f_{\texttt{joint}}$ and two
parallel branch layers $f_{\texttt{pre}}$ and $f_{\texttt{meta}}$. It enables
three interdependent data flows with different network parameters. In the
_joint-flow_ , the training data only go through $f_{\texttt{joint}}$, which
is jointly trained by pre-training and meta-training to extract common
knowledge as well as to coordinate the two parallel flows for a better
communication between them. Besides, the two parallel data flows for pre-
training and meta-training are respectively responsible for maintaining domain
transferability and task transferability. For pre-training, the data pass
through both $f_{\texttt{joint}}$ and $f_{\texttt{pre}}$, and then these two
outputs are added as the output of this Omni-Layer in the data flow. We denote
this data flow as _pre-flow_. Similarly, for meta-training and its
corresponding _meta-flow_ , the output is derived by adding the outputs of
$f_{\texttt{joint}}$ and $f_{\texttt{meta}}$. Overall, the transformation
function of the three parallel data flows in the $l$-th Omni-Layer can be
summarized as:

$\mathbf{x}^{l}=\left\\{\begin{array}[]{lcl}f_{\texttt{joint}}^{l}(\mathbf{x}^{l-1})+f_{\texttt{pre}}^{l}(\mathbf{x}^{l-1})&&{\text{$\mathbf{x}\in$
pre-flow}}\\\ &&\\\
f_{\texttt{joint}}^{l}(\mathbf{x}^{l-1})&&{\text{$\mathbf{x}\in$ joint-
flow}}\\\ &&\\\
f_{\texttt{joint}}^{l}(\mathbf{x}^{l-1})+f_{\texttt{meta}}^{l}(\mathbf{x}^{l-1})&&{\text{$\mathbf{x}\in$
meta-flow}}\end{array}\right.$ (3)

This architecture can be transformed from the layers in existing backbones by
copying their original layers as the chunk layer $f_{\texttt{joint}}$ and
adding two similar branch layers $f_{\texttt{pre}}$ and $f_{\texttt{meta}}$.
We design the two parallel branches as _lightweight_ layers compared to the
chunk layer, which maintains parameter efficiency of the Omni-Training
framework. For example, if $f_{\texttt{joint}}$ is a convolution layer with
large kernels such as $7\times 7$ or $3\times 3$, $f_{\texttt{pre}}$ and
$f_{\texttt{meta}}$ can be convolution layers with smaller kernels such as
$1\times 1$. Some existing architectures may introduce some additional special
layers such as batch normalization and various activation functions. We let
each data flow have its specific copy of these additional layers (denoted as
$a_{\texttt{joint}}$, $a_{\texttt{pre}}$ and $a_{\texttt{meta}}$), which
strengthens the specificity of the three data flows. We omit these additional
layers in the equations for simplicity.

We stack the Omni-Layers to construct the backbone for Omni-Training, and the
tri-flow in each layer expands to the entire data flows in the whole backbone.
Specifically, we use $F_{\texttt{joint}}$ to denote the overall function of
the _joint-flow_ which stacks $f_{\texttt{joint}}^{l}$ in the backbone:

$F_{\texttt{joint}}=f_{\texttt{joint}}^{L}\circ\cdots\circ
f_{\texttt{joint}}^{l}\circ\cdots\circ f_{\texttt{joint}}^{1}.$ (4)

We use $F_{\texttt{pre}}$ to denote the overall function of the stacked layers
in the backbone that encodes the _pre-flow_ , which enables knowledge routing
by adding the joint-flow:

$F_{\texttt{pre}}=\big{(}f_{\texttt{pre}}^{L}+f_{\texttt{joint}}^{L}\big{)}\circ\cdots\circ\big{(}f_{\texttt{pre}}^{l}+f_{\texttt{joint}}^{l}\big{)}\circ\cdots\circ\big{(}f_{\texttt{pre}}^{1}+f_{\texttt{joint}}^{1}\big{)}.$
(5)

Similarly, we use $F_{\texttt{meta}}$ to denote the overall function of the
stacked layers in the backbone that encodes the _meta-flow_ , which enables
knowledge routing by adding the joint-flow:

$\scriptsize
F_{\texttt{meta}}=\big{(}f_{\texttt{meta}}^{L}+f_{\texttt{joint}}^{L}\big{)}\circ\cdots\circ\big{(}f_{\texttt{meta}}^{l}+f_{\texttt{joint}}^{l}\big{)}\circ\cdots\circ\big{(}f_{\texttt{meta}}^{1}+f_{\texttt{joint}}^{1}\big{)}.$
(6)

Such a stacked tri-flow encoding backbone has several benefits. First, it is
parameter efficient, where the main chunk parameters are reused to encode
different data flows and the architecture requires much fewer parameters than
encoding these flows separately. Second, knowledge is softly shared between
pre-training, meta-training, and joint-training by routing through the shared
parameters in the architecture. Third, the Omni-Layer does not restrict on any
specific architecture choices, but is generally applicable to various
backbones in representation learning methods.

Omni-Head. The Omni-Head $H$ generates the final predictions of the three data
flows with the backbone representations. Specifically, $H$ consists of three
heads: a joint-head $H_{\texttt{joint}}$, a pre-head $H_{\texttt{pre}}$ and a
meta-head $H_{\texttt{meta}}$. Each head takes the corresponding data flow
representations in the backbone as its input and outputs the prediction.
Architectures of the three heads rely on the task, _e.g._ , for classification
problem, the heads can be classifiers with a single fully-connected layer. The
separate outputs for the three data flows enable the use of different losses
to train the three flows as introduced in Omni-Loss below. By chaining the
backbone and the Omni-Head, we obtain the Omni-Net architecture.

###  4.2 Omni-Loss

Based on the Omni-Net architecture, our general idea is to train the
parameters of each data flow with the corresponding pre-training or meta-
training algorithm, and enhance the transferability of each flow through the
Omni-Loss.

Joint-Training. Joint-training is performed on the joint-flow with the losses
of both pre-training and meta-training. In each iteration, we sample a
standard _mini-batch_ $\mathcal{B}$ and a _task episode_
$\\{\mathcal{S},\mathcal{Q}\\}$ from the large-scale training set
$\mathcal{D}_{\texttt{train}}$. We add the pre-training loss with the mini-
batch data and the meta-training loss with the sampled task on the joint-head
$H_{\texttt{joint}}$. The joint-training loss is

$\displaystyle\mathcal{L}_{\texttt{joint}}$ $\displaystyle=\
\mathbb{E}_{\mathcal{B}\sim\mathcal{P}(\mathbf{x},\mathbf{y})}\mathbb{E}_{(\mathbf{x},\mathbf{y})\in\mathcal{B}}\
\ell_{\texttt{pre}}\big{(}\mathbf{y},H_{\texttt{joint}}\circ
F_{\texttt{joint}}(\mathbf{x})\big{)}$ (7) $\displaystyle+\
\mathbb{E}_{(\mathcal{S},\mathcal{Q})\sim\mathcal{P}(\mathcal{T})}\mathbb{E}_{(\mathbf{x},\mathbf{y})\in\mathcal{Q}}\
\ell_{\texttt{meta}}\big{(}\mathbf{y},H_{\texttt{joint}}\circ
F_{\texttt{joint}}(\mathbf{x}|\mathcal{S})\big{)},$

where $\ell_{\texttt{pre}}$ and $\ell_{\texttt{meta}}$ are the losses of pre-
training and meta-training algorithms respectively. Though the joint-training
extracts shared features between the two training paradigms, such a naive
combination fails to endow representations with both domain transferability
and task transferability simultaneously, as we have shown in Section 3.3.
Therefore, we further perform pre-training and meta-training on the two
parallel data flows respectively to explicitly preserve domain transferability
and task transferability.

Pre-Training. To specifically acquire domain transferability in the network,
we perform pre-training on the pre-flow. In each iteration, we feed each
sample $(\mathbf{x},\mathbf{y})$ from the mini-batch $\mathcal{B}$ into the
pre-flow of the Omni-Net, going through $F_{\texttt{pre}}$ and
$H_{\texttt{pre}}$, and control the final output by the pre-training loss on
the pre-flow:

$\mathcal{L}_{\texttt{pre}}=\mathbb{E}_{\mathcal{B}\sim\mathcal{P}(\mathbf{x},\mathbf{y})}\mathbb{E}_{(\mathbf{x},\mathbf{y})\in\mathcal{B}}\
\ell_{\texttt{pre}}\big{(}\mathbf{y},H_{\texttt{pre}}\circ
F_{\texttt{pre}}(\mathbf{x})\big{)}.$ (8)

In addition to the knowledge transfer across different branches, we further
enhance the specific transferability on each parallel branch throughout the
learning process. In order to realize it, we employ a self-distillation
strategy. Let $\theta$ denote all the parameters in the backbone $F$ and the
Omni-Head $H$, $i$ denote the training steps, we keep the temporal ensemble of
the network during the learning process, _i.e._ , an exponential moving
average (EMA) of the model parameters $\widetilde{\theta}$, which is updated
smoothly during training:

$\widetilde{\theta}_{i}=\alpha\widetilde{\theta}_{i-1}+(1-\alpha)\theta_{i}.$
(9)

The EMA model gathers knowledge from different training stages and serves as a
teacher to guide the training of the current Omni-Net. In each iteration, the
EMA model transfers knowledge to each parallel branch through knowledge
distillation. We implement this idea into _self-distillation regularization_
for the pre-flow:

$\mathcal{R}_{\texttt{pre}}=\mathbb{E}_{\mathcal{B}\sim\mathcal{P}(\mathbf{x},\mathbf{y})}\mathbb{E}_{(\mathbf{x},\mathbf{y})\in\mathcal{B}}\
\ell_{2}\big{(}\widetilde{H}_{\texttt{pre}}\circ\widetilde{F}_{\texttt{pre}}(\mathbf{x}),H_{\texttt{pre}}\circ
F_{\texttt{pre}}(\mathbf{x})\big{)},$ (10)

where $\widetilde{F}_{\texttt{pre}}$ and $\widetilde{H}_{\texttt{pre}}$ denote
the mapping functions of pre-flow and pre-head in the EMA model with the
temporal ensemble parameters of $\widetilde{\theta}$, and $\ell_{2}$ is the
squared loss. The pre-training loss improved by the self-distillation for the
pre-flow is

$\mathcal{J}_{\texttt{pre}}=\mathcal{L}_{\texttt{pre}}+\lambda\mathcal{R}_{\texttt{pre}},$
(11)

with $\lambda$ being a hyper-parameter to trade-off the original pre-training
loss and the self-distillation regularization.

Meta-Training. To acquire task transferability in the network, in each
iteration, we perform meta-training on the meta-flow with the sampled task
episode $(\mathcal{S},\mathcal{Q})$. Data in the support set $\mathcal{S}$ are
fed into the meta-flow to obtain the conditioned model. Then, each sample
$(\mathbf{x},\mathbf{y})$ from the query set $\mathcal{Q}$ passes through the
meta-flow conditioned on the support set to derive the meta-training loss:

$\mathcal{L}_{\texttt{meta}}=\mathbb{E}_{(\mathcal{S},\mathcal{Q})\sim\mathcal{P}(\mathcal{T})}\mathbb{E}_{(\mathbf{x},\mathbf{y})\in\mathcal{Q}}\
\ell_{\texttt{meta}}\big{(}\mathbf{y},H_{\texttt{meta}}\circ
F_{\texttt{meta}}(\mathbf{x}|\mathcal{S})\big{)}.$ (12)

Similar to the pre-flow, we impose the _self-distillation regularization_ to
improve the transferability of the meta-learned representations across the
training process for the meta-flow:

$\displaystyle\mathcal{R}_{\texttt{meta}}=\mathbb{E}_{(\mathcal{S},\mathcal{Q})\sim\mathcal{P}(\mathcal{T})}\mathbb{E}_{(\mathbf{x},\mathbf{y})\in\mathcal{Q}}\
\ell_{2}\big{(}\widetilde{H}_{\texttt{meta}}\circ\widetilde{F}_{\texttt{meta}}(\mathbf{x}|\mathcal{S}),$
(13) $\displaystyle H_{\texttt{meta}}\circ
F_{\texttt{meta}}(\mathbf{x}|\mathcal{S})\big{)},$

where $\widetilde{F}_{\texttt{meta}}$ and $\widetilde{H}_{\texttt{meta}}$
denote the mapping functions of the meta-flow and meta-head in the EMA model,
and $\ell_{2}$ is the squared loss. The training loss for the meta-flow
includes the original meta-training loss and the self-distillation
regularization as

$\mathcal{J}_{\texttt{meta}}=\mathcal{L}_{\texttt{meta}}+\lambda\mathcal{R}_{\texttt{meta}},$
(14)

with $\lambda$ to trade-off the original meta-training loss and the
regularization term.

###  4.3 Overall Framework

Training. We train Omni-Net with the Omni-Loss to perform joint-training, pre-
training and meta-training simultaneously:

$\mathcal{O}_{\texttt{Omni}}=\mathcal{J}_{\texttt{pre}}+\mathcal{J}_{\texttt{meta}}+\mathcal{L}_{\texttt{joint}}.$
(15)

With the cooperation of Omni-Net and Omni-Loss, our framework trains the two
parallel flows to obtain both domain transferability and task transferability
and coordinates the two parallel flows to enable their knowledge
communication, addressing both challenges of _domain shift_ and _task shift_
in few-shot learning problems.

Inference. During the test time, we transfer knowledge learned from Omni-
Training by reusing or fine-tuning the learned model and retraining a new
Omni-Head for the new tasks on the labeled data in the support set
$\mathcal{S}_{\texttt{test}}$. Since we focus on the representation learning
stage but do not focus on the test time adaptation techniques, we train the
new Omni-Head consisting of a new joint-head
$H^{\texttt{new}}_{\texttt{joint}}$, a new pre-head
$H^{\texttt{new}}_{\texttt{pre}}$ and a new meta-head
$H^{\texttt{new}}_{\texttt{meta}}$ following the corresponding algorithms we
have used for pre-training and meta-training. Then for each test sample
$\mathbf{x}\in\mathcal{Q}_{\texttt{test}}$, we predict $\mathbf{x}$ using one
of the three heads or their ensemble based on the real application
constraints. For example, if we need to deploy the model to a real-time
prediction application, we only use the prediction of the meta-head for fast
adaptation using only a few gradient updates. If there is no resource
restriction, we can use the ensemble of all three heads for more accurate
predictions.

##  5 Omni-Training Algorithms

We provide instantiations and implementations of the Omni-Training framework
by incorporating some mainstream pre-training and meta-training algorithms.
The framework can generalize to a wider variety of algorithms as shown in our
experiments.

###  5.1 Pre-Training Algorithms

Classification. The pre-training algorithm for classification is known as
Baseline [15] in few-shot learning literature. To instantiate,
$H_{\texttt{pre}}$ is a fully-connected layer with weights
$[\mathbf{w}_{1},...,\mathbf{w}_{K}]$ and biases $[b_{1},...,b_{K}]$ for $K$
classes, $F_{\texttt{pre}}$ and $H_{\texttt{pre}}$ are pre-trained on training
dataset $\mathcal{D}_{\texttt{train}}$ by using cross-entropy as
$\ell_{\texttt{pre}}$:

$\ell_{\texttt{pre}}\big{(}\mathbf{y},H_{\texttt{pre}}\circ
F_{\texttt{pre}}(\mathbf{x})\big{)}=-\log\left(\frac{\exp(\mathbf{w}_{y}^{T}F_{\texttt{pre}}(\mathbf{x}))}{\sum_{k}\exp(\mathbf{w}_{k}^{T}F_{\texttt{pre}}(\mathbf{x}))}\right),$
(16)

where $y$ is the class index of the ground-truth class label $\mathbf{y}$ for
$\mathbf{x}$. The model is then fine-tuned on the support set
$\mathcal{S}_{\texttt{test}}$ for the new task with a new classification head
$H^{\texttt{new}}_{\texttt{pre}}$.

Regression. In the pre-training algorithm for regression, we use a fully-
connected layer as the pre-head $H_{\texttt{pre}}$ to predict the output. Here
the loss is defined as the squared error between the target value $y$ and the
prediction, also known as the L2 loss:

$\ell_{\texttt{pre}}\big{(}y,H_{\texttt{pre}}\circ
F_{\texttt{pre}}(\mathbf{x})\big{)}=\big{(}H_{\texttt{pre}}\circ
F_{\texttt{pre}}(\mathbf{x})-y\big{)}^{2}.$ (17)

Reinforcement Learning. In the pre-training algorithm for reinforcement
learning, we use the policy gradient in REINFORCE [75]. The Omni-Net serves as
the parameterized policy $\pi=H_{\texttt{pre}}\circ F_{\texttt{pre}}$ with a
fully-connected head $H_{\texttt{pre}}$ to predict the action given a state.
Here the loss is defined as the expected return over the policy:
$\ell_{\texttt{pre}}\left(H_{\texttt{pre}}\circ
F_{\texttt{pre}}\right)=\mathbb{E}_{\tau\sim\pi}\big{[}\sum\nolimits_{t=0}^{\infty}{r}(s_{t},a_{t})\big{]}$.
The gradient of the pre-training loss $\ell_{\texttt{pre}}$ with respect to
the parameters $\theta$ of the policy $\pi$, _i.e._ , the policy gradient, is
defined as

$\nabla_{\theta}\ell_{\texttt{pre}}\left(H_{\texttt{pre}}\circ
F_{\texttt{pre}}\right)=\sum\limits_{s}p^{\pi}(s)\sum\limits_{a}\nabla_{\theta}\pi(a|s)Q^{\pi}(s,a).$
(18)

$p^{\pi}(s)$ is discounted weighting of the probability of encountering states
$s$ from the initial states and $Q^{\pi}$ is the Q-function for $\pi$ [76].

###  5.2 Meta-Training Algorithms

Model-Agnostic Meta-Learning (MAML). In meta-training, we first consider
model-agnostic meta-learning (MAML) [9], a gradient-based learning rule to
rapidly adapt to new tasks with few data and gradient steps. In each
iteration, we sample an episode of a support set $\mathcal{S}$ and a query set
$\mathcal{Q}$, and optimize the MAML loss:

$\displaystyle\ell_{\texttt{meta}}\big{(}\mathbf{y},H_{\texttt{meta}}\circ
F_{\texttt{meta}}(\mathbf{x}|\mathcal{S})\big{)}=\ell\big{(}\mathbf{y},H_{\texttt{meta}}\circ
F_{\texttt{meta}}(\mathbf{x};\theta^{\prime})\big{)},$ (19)

for each sample $(\mathbf{x},\mathbf{y})\in\mathcal{Q}$ in the query set. Here
$\theta$ is the parameters of $H_{\texttt{meta}}$ and $F_{\texttt{meta}}$ in
the meta-flow, and
$\theta^{\prime}=\theta-\nabla_{\theta}\mathbb{E}_{(\mathbf{x},\mathbf{y})\in\mathcal{S}}\
\ell\big{(}\mathbf{y},H_{\texttt{meta}}\circ
F_{\texttt{meta}}(\mathbf{x};\theta)\big{)}$ is the model parameters after a
single gradient update on the support set $\mathcal{S}$. MAML has few
restrictions on the model architecture and learning task, and can be widely
used on various tasks such as regression, classification and reinforcement
learning, by specifying the task-aware loss $\ell$.

Prototypical Networks. In the few-shot learning literature, one of the well-
established meta-training algorithms is ProtoNet [14]. Let $\mathcal{S}_{k}$
denote the samples with the class index $k$ in a support set $\mathcal{S}$ in
the episode, the prototype of this class $\mathbf{c}_{k}$ is the mean of the
embedded data in $\mathcal{S}_{k}$:
$\mathbf{c}_{k}=\mathbb{E}_{(\mathbf{x},\mathbf{y})\in{\mathcal{S}_{k}}}F_{\texttt{meta}}(\mathbf{x})$.
A metric-based classifier predicts the probability distribution of each query
point $\mathbf{x}$ based on its Euclidean distances $d$ to the prototypes,
which is penalized by a cross-entropy loss for classification:

$\scriptsize\ell_{\texttt{meta}}\big{(}\mathbf{y},H_{\texttt{meta}}\circ
F_{\texttt{meta}}(\mathbf{x}|\mathcal{S})\big{)}=-\log\left(\frac{\exp(-d(F_{\texttt{meta}}(\mathbf{x}),\mathbf{c}_{y}))}{\sum_{k=1}^{K}\exp(-d(F_{\texttt{meta}}(\mathbf{x}),\mathbf{c}_{k}))}\right).$
(20)

For new tasks, the labeled data in the support set
$\mathcal{S}_{\texttt{test}}$ are used to compute the prototypes of each new
class. Then we can classify new samples in the query set
$\mathcal{Q}_{\texttt{test}}$ by their nearest prototype.

##  6 Experiments

We evaluate our Omni-Training framework with comprehensive experiments on
cross-task and cross-domain settings in classification, regression and
reinforcement learning problems to testify the few-shot learning performances.
All the codes and datasets will be available online at
https://github.com/thuml/Omni-Training.

###  6.1 Classification

TABLE I: The results of the new tasks with $5$ or $1$ labeled samples per class on mini-ImageNet and CUB datasets. Method | Backbone | mini-ImageNet | CUB
---|---|---|---
$K=5$ | $K=1$ | $K=5$ | $K=1$
MatchingNet [13] | ResNet-18 | $68.88\pm 0.69$ | $52.91\pm 0.88$ | $83.64\pm 0.60$ | $72.36\pm 0.90$
ProtoNet [14] | ResNet-18 | $73.68\pm 0.65$ | $54.16\pm 0.82$ | $87.42\pm 0.48$ | $71.88\pm 0.91$
RelationNet [42] | ResNet-18 | $69.83\pm 0.68$ | $52.48\pm 0.86$ | $82.75\pm 0.58$ | $67.59\pm 1.02$
MAML [9] | ResNet-18 | $65.72\pm 0.77$ | $49.61\pm 0.92$ | $82.70\pm 0.65$ | $69.96\pm 1.01$
TADAM [46] | ResNet-12 | $76.70\pm 0.30$ | $58.50\pm 0.30$ | $-$ | $-$
GNN [44] | ResNet-18 | $78.80\pm 0.78$ | $57.40\pm 0.98$ | $90.74\pm 0.57$ | $78.52\pm 1.03$
LEO [69] | WRN28-10 | $77.59\pm 0.12$ | $61.76\pm 0.08$ | $78.27\pm 0.16$ | $68.22\pm 0.22$
Baseline [15] | ResNet-18 | $74.27\pm 0.63$ | $51.75\pm 0.80$ | $82.85\pm 0.55$ | $65.51\pm 0.87$
Baseline++ [15] | ResNet-18 | $75.68\pm 0.63$ | $51.87\pm 0.77$ | $83.58\pm 0.54$ | $67.02\pm 0.90$
MTL [73] | ResNet-12 | $75.50\pm 0.80$ | $61.20\pm 1.80$ | $-$ | $-$
MetaOpt [71] | ResNet-12 | $78.63\pm 0.36$ | $62.64\pm 0.61$ | $90.90\pm 0.23$ | $80.23\pm 0.44$
TapNet [47] | ResNet-12 | $76.36\pm 0.10$ | $61.65\pm 0.15$ | $-$ | $-$
Robust20 [29] | ResNet-18 | $81.59\pm 0.42$ | $63.95\pm 0.61$ | $84.62\pm 0.44$ | $69.47\pm 0.69$
CAN [45] | ResNet-12 | $79.44\pm 0.34$ | $63.85\pm 0.48$ | $-$ | $-$
RFS [30] | ResNet-12 | $82.14\pm 0.43$ | $64.82\pm 0.60$ | $-$ | $-$
Neg-Margin [77] | ResNet-18 | $80.94\pm 0.59$ | $62.33\pm 0.82$ | $89.40\pm 0.40$ | $72.66\pm 0.90$
PMM [78] | ResNet-18 | $77.76\pm 0.58$ | $60.11\pm 0.73$ | $86.01\pm 0.50$ | $73.94\pm 1.10$
Multi-Task [79] | ResNet-12 | $77.72\pm 0.09$ | $59.84\pm 0.22$ | $-$ | $-$
Meta-Maxup [80] | ResNet-12 | $79.38\pm 0.24$ | $62.81\pm 0.34$ | $-$ | $-$
OT-Proto | ResNet-18 | $81.26\pm 0.57$ | $64.31\pm 0.86$ | $91.09\pm 0.38$ | $81.18\pm 0.78$
OT-Proto | ResNet-12 | ${82.36\pm 0.54}$ | ${66.62\pm 0.80}$ | ${91.93\pm 0.38}$ | ${82.94\pm 0.73}$
OT-GNN | ResNet-18 | $\mathbf{87.14\pm 0.59}$ | $\mathbf{70.99\pm 0.97}$ | $\mathbf{95.96\pm 0.33}$ | $\mathbf{87.73\pm 0.78}$

Datasets. We consider few-shot classification problems with four datasets: in-
domain datasets mini-ImageNet [13], CUB [81], and cross-domain datasets mini-
ImageNet$\rightarrow$CUB, Multi-domain. mini-ImageNet is a subset of ILSVRC-12
dataset [82] for generic object recognition. It contains $100$ classes with
$600$ images per class. We use the same split introduced by [52], which
respectively splits $64$/$16$/$20$ classes for the training/validation/testing
set. CUB is a fine-grained dataset of birds with a total of $200$ classes and
$11,788$ images. We follow the protocol of [83] and split the dataset into
$100$/$50$/$50$ classes for training/validation/testing. mini-
ImageNet$\rightarrow$CUB is a cross-domain dataset. Following [15], we use
mini-ImageNet dataset as the training set and split the CUB set as $50$/$50$
classes for validation and testing. Multi-domain is another cross-domain
dataset. We follow the split in [72] and use the datasets of mini-ImageNet,
CUB, Cars [84], Places [85] and Plantae [86] as different domains. We explore
two settings. The first is training the model on the mini-ImageNet domain and
evaluating on other four domains. The second is the leave-one-out setting
which selects one domain for evaluation and trains the model with all other
domains.

TABLE II: The classification accuracy of the new tasks with $5$ or $1$ labeled samples per class in the _cross-domain_ setting, mini-ImageNet$\rightarrow$CUB. Method | $K=5$ | $K=1$
---|---|---
MatchingNet [13] | $53.07\pm 0.74$ | $38.78\pm 0.73$
ProtoNet [14] | $62.02\pm 0.70$ | $40.07\pm 0.75$
RelationNet [42] | $57.71\pm 0.73$ | $37.71\pm 0.69$
MAML [9] | $51.34\pm 0.72$ | $40.15\pm 0.65$
GNN [44] | $65.56\pm 0.87$ | $43.65\pm 0.86$
Baseline [15] | $65.57\pm 0.70$ | $43.59\pm 0.74$
Baseline++ [15] | $62.04\pm 0.76$ | $44.14\pm 0.77$
Robust20 [29] | $65.04\pm 0.57$ | $-$
Neg-Margin [77] | $67.03\pm 0.76$ | $-$
PMM [78] | $68.77\pm 0.90$ | $-$
OT-Proto | $71.30\pm 0.71$ | $50.42\pm 0.82$
OT-GNN | $\mathbf{75.83\pm 0.82}$ | ${50.89\pm 0.91}$

Implementation Details. We use ResNet-18 in [15] and ResNet-12 with dropblocks
in [46] as the backbone for mini-ImageNet, CUB and mini-
ImageNet$\rightarrow$CUB. Following [72], we use ResNet-10 on Multi-domain for
a fair comparison. We refactor ResNet into a backbone for Omni-Training by
transforming all convolution layers into Omni-Layers, where each Omni-Layer
uses the $1\times 1$ convolution layer as the lightweight branch layer. We
employ Baseline in [15] as the pre-training method and explore two powerful
meta-training methods, ProtoNet [14] and GNN [44], denoted as OT-Proto and OT-
GNN respectively. In each iteration, a mini-batch is sampled with the batch
size of $64$ for pre-training, and an episode of task is sampled for meta-
training, with a support set containing $5$ categories each having $5$ labeled
instances, and a query set containing the same categories with $16$ instances
per class. We apply standard data augmentation including random crop, left-
right flip and color jitter to the training samples. We train our framework
with $100$ epochs for the mini-ImageNet, mini-ImageNet$\rightarrow$CUB and
Multi-domain datasets, and with $400$ epochs for the CUB dataset. We use
accuracy on the validation set to choose the best model for testing. In the
test stage, we randomly sample $600$ tasks from the testing set. Each task
contains $5$ unseen classes with $K=5$ or $K=1$ labeled samples per class as
the support set, and another $16$ instances per class as the query set to be
predicted. The average accuracy as well as the $95\%$ confidence intervals are
reported. The hyper-parameter is chosen as $\alpha=0.99$, $\lambda=3.0$. We
train the networks from scratch and use Adam optimizer [87] with an initial
learning rate of $0.001$.

TABLE III: The results of the tasks from unseen domains with $5$ or $1$
labeled samples per class in the _Multi-domain_ setting (trained with mini-
ImageNet).

Method | CUB | Cars | Places | Plantae
---|---|---|---|---
$K=5$ | $K=1$ | $K=5$ | $K=1$ | $K=5$ | $K=1$ | $K=5$ | $K=1$
MatchingNet [13] | $51.37\pm 0.77$ | $35.89\pm 0.51$ | $38.99\pm 0.64$ | $30.77\pm 0.47$ | $63.16\pm 0.77$ | $49.86\pm 0.79$ | $46.53\pm 0.68$ | $32.70\pm 0.60$
ProtoNet [14] | $57.64\pm 0.85$ | $38.18\pm 0.76$ | $42.84\pm 0.73$ | $29.72\pm 0.59$ | $68.86\pm 0.70$ | $49.24\pm 0.81$ | $47.41\pm 0.70$ | $35.02\pm 0.63$
RelationNet [42] | $57.77\pm 0.69$ | $42.44\pm 0.77$ | $37.33\pm 0.68$ | $29.11\pm 0.60$ | $63.32\pm 0.76$ | $48.64\pm 0.85$ | $44.00\pm 0.60$ | $33.17\pm 0.64$
GNN [44] | $62.25\pm 0.65$ | $45.69\pm 0.68$ | $44.28\pm 0.63$ | $31.79\pm 0.51$ | $70.84\pm 0.65$ | $53.10\pm 0.80$ | $52.53\pm 0.59$ | $35.60\pm 0.56$
FT-Matching [72] | $55.23\pm 0.83$ | $36.61\pm 0.53$ | $41.24\pm 0.65$ | $29.82\pm 0.44$ | $64.55\pm 0.75$ | $51.07\pm 0.68$ | $41.69\pm 0.63$ | $34.48\pm 0.50$
FT-Relation [72] | $59.46\pm 0.71$ | $44.07\pm 0.77$ | $39.91\pm 0.69$ | $28.63\pm 0.59$ | $66.28\pm 0.72$ | $50.68\pm 0.87$ | $45.08\pm 0.59$ | $33.14\pm 0.62$
FT-GNN [72] | ${66.98\pm 0.68}$ | ${47.47\pm 0.75}$ | $44.90\pm 0.64$ | $31.61\pm 0.53$ | $73.94\pm 0.67$ | ${55.77\pm 0.79}$ | $53.85\pm 0.62$ | $35.95\pm 0.58$
OT-Proto | $65.17\pm 0.75$ | $45.83\pm 0.78$ | $\mathbf{51.19\pm 0.71}$ | ${34.82\pm 0.70}$ | ${74.16\pm 0.69}$ | $55.73\pm 0.89$ | ${57.88\pm 0.69}$ | ${39.51\pm 0.71}$
OT-GNN | $\mathbf{70.24\pm 0.82}$ | $\mathbf{49.51\pm 0.96}$ | ${48.99\pm 0.83}$ | $\mathbf{35.31\pm 0.78}$ | $\mathbf{79.61\pm 0.81}$ | $\mathbf{61.74\pm 1.05}$ | $\mathbf{60.08\pm 0.81}$ | $\mathbf{40.52\pm 0.81}$

TABLE IV: The results of the tasks from unseen domains with $5$ or $1$ labeled
samples per class in the _Multi-domain_ setting (trained with all other
domains).

Method | CUB | Cars | Places | Plantae
---|---|---|---|---
$K=5$ | $K=1$ | $K=5$ | $K=1$ | $K=5$ | $K=1$ | $K=5$ | $K=1$
MatchingNet [13] | $51.92\pm 0.80$ | $37.90\pm 0.55$ | $39.87\pm 0.51$ | $28.96\pm 0.45$ | $61.82\pm 0.57$ | $49.01\pm 0.65$ | $47.29\pm 0.51$ | $33.21\pm 0.51$
ProtoNet [14] | $59.26\pm 0.89$ | $39.31\pm 0.72$ | $43.66\pm 0.68$ | $29.52\pm 0.54$ | $68.03\pm 0.61$ | $47.96\pm 0.77$ | $49.35\pm 0.72$ | $35.40\pm 0.68$
RelationNet [42] | $62.13\pm 0.74$ | $44.33\pm 0.59$ | $40.64\pm 0.54$ | $29.53\pm 0.45$ | $64.34\pm 0.57$ | $47.76\pm 0.63$ | $46.29\pm 0.56$ | $33.76\pm 0.52$
GNN [44] | $69.26\pm 0.68$ | $49.46\pm 0.73$ | $48.91\pm 0.67$ | $32.95\pm 0.56$ | $72.59\pm 0.67$ | $51.39\pm 0.80$ | $58.36\pm 0.68$ | $37.15\pm 0.60$
FT-Matching [72] | $61.41\pm 0.57$ | $43.29\pm 0.59$ | $43.08\pm 0.55$ | $30.62\pm 0.48$ | $64.99\pm 0.59$ | $52.51\pm 0.67$ | $48.32\pm 0.57$ | $35.12\pm 0.54$
FT-Relation [72] | $64.99\pm 0.54$ | $48.38\pm 0.63$ | $43.44\pm 0.59$ | $32.21\pm 0.51$ | $67.35\pm 0.54$ | $50.74\pm 0.66$ | $50.39\pm 0.52$ | $35.00\pm 0.52$
FT-GNN [72] | $\mathbf{73.11\pm 0.68}$ | $\mathbf{51.51\pm 0.80}$ | $49.88\pm 0.67$ | $34.12\pm 0.63$ | $77.05\pm 0.65$ | ${56.31\pm 0.80}$ | $58.84\pm 0.66$ | $42.09\pm 0.68$
OT-Proto | $67.76\pm 0.74$ | $46.62\pm 0.77$ | ${52.02\pm 0.74}$ | ${36.36\pm 0.70}$ | ${73.57\pm 0.66}$ | $52.20\pm 0.81$ | ${59.37\pm 0.69}$ | ${40.95\pm 0.66}$
OT-GNN | $71.78\pm 0.83$ | $49.78\pm 0.94$ | $\mathbf{54.39\pm 0.83}$ | $\mathbf{37.00\pm 0.79}$ | $\mathbf{78.03\pm 0.80}$ | $\mathbf{58.27\pm 0.99}$ | $\mathbf{62.22\pm 0.79}$ | $\mathbf{43.02\pm 0.85}$

Results on Cross-Task Benchmarks. We first evaluate our method on the general
dataset mini-ImageNet and the fine-grained dataset CUB. These two scenarios
are considered as _cross-task_ benchmarks, as the training and testing data
are from the same domain. The results with $K=5$ and $K=1$ are shown in Table
I. Omni-Training outperforms corresponding pre-training and meta-training
methods, especially in extremely difficult scenarios with only $1$ labeled
instance. Note that although from the same dataset, there still exists domain
shift between the training and test sets caused by the split of different
label sets. Our framework manages to incorporate pre-training and meta-
training effectively to acquire both domain transferability and task
transferability and thus achieves higher performance. Omni-Training
outperforms state-of-the-art algorithms, including MTL [73] which combines
pre-training and meta-training sequentially. This confirms that our design can
better bridge pre-training and meta-training.

Results on Cross-Domain Benchmarks. We consider two more challenging _cross-
domain_ benchmarks, mini-imageNet$\rightarrow$CUB and Multi-domain. Different
from the cross-task benchmarks discussed above, in the cross-domain setting,
the testing data are not only from different classes, but also from different
domains, causing greater domain shift between the training and testing data.
As shown in Table II, meta-training algorithms degrade due to the domain shift
while pre-training algorithms generalize better to the unseen domain. Omni-
Training outperforms meta-training methods by a large margin, indicating the
significance of domain transferability in the cross-domain setting. Also,
Omni-Training outperforms the pre-training Baseline, which reveals the
importance of task transferability to fully enable few-shot learning.

A more challenging benchmark is Multi-domain with more domains and larger
domain shift. Table III reports the results of the first setting where we
train on the mini-ImageNet domain and test on other four domains. Table IV
reports the results of the second leave-one-out setting, where we choose one
domain as the unseen test domain and train the model with all other domains.
We specially compare Omni-Training with Feature-Transformation [72], which is
a framework adopts domain generalization [88] into meta-training to obtain
both domain transferability and task transferability. Among its
implementations, FT-GNN achieves the best performance by incorporating a
strong meta-training algorithm, GNN [44]. When trained on mini-ImageNet, Omni-
Training can still achieve comparable or better performance than FT-GNN with a
simple meta-training algorithm such as ProtoNet. We also incorporate GNN into
Omni-Training to form OT-GNN, which generally outperforms FT-GNN in most
tasks. Note that FT-GNN has a special design for domain generalization, which
is tailored for the multi-domain setting. But OT-GNN also achieves better
performance on most cases, confirming that Omni-Training works generally well
in different situations.

###  6.2 Regression

Datasets. For few-shot regression problems, we conduct experiments on a
sinusoid dataset following [9]. Specifically, the regression problem is to
predict the output $y$ on a sine wave given the input $x$. We define a task as
regressing a sine wave with a particular amplitude and phase from some labeled
data and consider a continuous task distribution in which the amplitude varies
within $[0.1,5.0]$ and the phase varies within $[0,2\pi]$. The input datapoint
$x$ is sampled uniformly from $[-5.0,5.0]$ for all tasks. The training dataset
$\mathcal{D}_{\texttt{train}}$ contains a large number of sampled sine waves
and each test task
$\\{\mathcal{S}_{\texttt{test}},\mathcal{Q}_{\texttt{test}}\\}$ is an unseen
sinusoid with a few labeled datapoints in $\mathcal{S}_{\texttt{test}}$ and
other points which need prediction in $\mathcal{Q}_{\texttt{test}}$. The goal
is to train a regression model on $\mathcal{D}_{\texttt{train}}$ to predict
the outputs of the datapoints in the query set $\mathcal{Q}_{\texttt{test}}$
after adaptation with a few labeled data in $\mathcal{S}_{\texttt{test}}$.

Implementation Details. We take the mean-squared error between the predictions
and ground-truth values as the training loss. We use Baseline [15] for pre-
training and use MAML [9] for meta-training. We employ a backbone with $2$
fully-connected layers of size $64$ with the activation function of Tanh. The
training set $\mathcal{D}_{\texttt{train}}$ has 30000 randomly sampled tasks
and each task is a sine wave with $50$ labeled datapoints. We then enable few-
shot regression on a new sine wave with a support set of $K=\\{5,10,20\\}$
labeled examples and test the adapted model on points in
$\mathcal{Q}_{\texttt{test}}$ of the wave. We train the model on
$\mathcal{D}_{\texttt{train}}$ and fine-tune it on the $K$ labeled examples
for the new sine wave with an SGD optimizer. The learning rate for the inner
loop is $0.02$ and that for parameter update is initialized as $0.01$.

(a) $K=5$

(b) $K=10$

(c) $K=20$

Figure 4: The few-shot regression results of different training methods with
different gradients steps and different support set sizes $K$.

(a) Baseline

(b) MAML

(c) Omni-Training

Figure 5: The recovered sine wave of Baseline, MAML and Omni-Training. The
models are updated using $5$ sampled points with $1$ or $10$ gradient steps.

Results. We sample 100 new tasks for testing and report the mean squared error
after fine-tuning with different gradient steps from $1$ to $10$. As shown in
Figure 4, Baseline generally performs worse than MAML. The tasks change
rapidly during the training and test stages in this problem and task
transferability is important, which is missing for pre-training methods. With
different numbers of labeled data and of gradient steps, Omni-Training
consistently improves upon the meta-training method, which shows the efficacy
of Omni-Training for regression tasks.

We further conduct a case study and show the typical sine waves recovered by
pre-training, meta-training and Omni-Training with $5$ labeled samples and
with $1$ or $10$ gradient steps in Figure 5. We also show the ground-truth
sine wave and the labeled points in the support set. MAML and Omni-Training
quickly regress closed to the ground-truth curve, while the process is much
slower for Baseline. Compared with MAML, the recovered curve of Baseline
maintains smooth, which is an important common property in the sinusoid
distribution. Omni-Training also maintains a smoother curve, which
simultaneously fits these datapoints quickly and preserves the domain
transferability of sine waves. This explains the improvements brought by the
Omni-Training framework.

(a) 2D Navigation

(b) Locomotion-Velocity

(c) Locomotion-Direction

Figure 6: Average expected return of the tasks in the 2D Navigation
environment and the Locomotion environment for reinforcement learning.

###  6.3 Reinforcement Learning

Environments. For reinforcement learning problems, we follow the learning
protocol in [9] with several sets of tasks based on two simulated continuous
control environments: 2D Navigation and Locomotion in the rllab benchmark
suite [89].

In the 2D Navigation environment, the goal is to move to a target position in
2D. The state space is the 2D location and the action space is the 2D
velocity, where the action is in the range of $\left[-0.1,0.1\right]$. The
reward is the negative squared distance to the goal, and the episodes
terminate when the agent is within $0.01$ of the goal or at the horizon of
$H=100$. We construct a task by randomly sampling a goal position from a unit
square.

In the Locomotion environment, we adopt the agent in the Mujoco HalfCheetah
environment [90] and follow its state and action space. We evaluate on two
sets of tasks. The first aims to run at a particular velocity. The reward is
the negative absolute value between the current velocity and the goal
velocity, which is chosen uniformly at random between $0.0$ and $2.0$ for
different tasks. The second aims to run in a particular direction. The reward
is the magnitude of the velocity in the forward or backward direction. The
horizons of both tasks are set as $H=200$.

Implementation Details. We adopt the policy as a neural network with two
fully-connected layers of $64$ hidden units and the Tanh activation function.
We train the policy with the REINFORCE algorithm [91]. We use the standard
linear feature baseline proposed by [89], which is fitted separately at each
iteration for each sampled task in the batch. We train the model with $500$
iterations. In each iteration, $20$ different tasks are sampled for the 2D
navigation environment and $40$ tasks are sampled for Locomotion, where $20$
trajectories are sampled for each task. During the test stage for few-shot
reinforcement learning, we randomly sample $2000$ new tasks for evaluation.
Each task contains trajectories with rewards as the support set. We use $20$
trajectories from each task for each gradient step and use $1$ to $4$ gradient
steps for adaptation to new tasks. We use $20$ trajectories as the query set
to compute the final testing reward of each task. We also use Baseline [15] as
the pre-training method and MAML [9] as the meta-training method. In each
iteration of meta-training, the policy is first trained using a single
gradient step on the support trajectories with the inner loop step size $0.1$,
and then meta-updated on the query trajectories with the outer loop step size
$0.03$.

Results. The results of the reinforcement learning tasks in the two
environments are shown in Figure 6. Omni-Training outperforms both Baseline
and MAML with large margins in the 2D Navigation environment, which
demonstrates that the model with both domain and task transferability can
boost the generalization performance in this case. In the Locomotion
environment, the performance gap between MAML and Baseline becomes larger,
indicating more complex cross-task situations. Omni-Training still improves
upon MAML in the velocity tasks. In the direction tasks, the pre-training
method fails to generalize across these complex tasks with limited
trajectories and updates, thereby performing similarly to the random
initialization. In this extreme case, Omni-Training still performs comparably
with MAML, without being negatively influenced. These results have proved the
generalization ability of Omni-Training in a variety of complex situations.

(a) mini-ImageNet$\rightarrow$CUB

(b) Regression

(c) Extension

Figure 7: (a)(b) Comparison with two ways of simply combining pre-training and
meta-training: Ensemble and Joint-Training. (c) Extension of the Omni-Training
framework to other representation learning algorithms (OT: Omni-Training).

(a) Pre-Training

(b) Meta-Training

(c) Transferability

Figure 8: (a)(b) Training losses and validation accuracies of the pre-flow and
meta-flow in Omni-Training, compared with their corresponding baselines; (c)
The transferability of pre-training and pre-flow with different numbers of
gradient steps.

##  7 Analysis

In this section, we further empirically analyze and understand our proposed
framework. Without specification, we use the ResNet-18 as the backbone. We use
the Baseline in [15] as the pre-training method and the ProtoNet [14] as the
meta-training method.

###  7.1 Fine-grained Comparison with Baselines

Comparison with Simple Combinations. We compare Omni-Training with two simple
combinations of pre-training and meta-training discussed in Section 3.3,
_i.e._ the ensemble of the two models trained separately (Ensemble) and joint-
training with the losses of the two training paradigms (Joint-Training). We
evaluate on the classification dataset mini-ImageNet$\rightarrow$CUB and the
sinusoid regression dataset. We use $K=5$ or $K=1$ labeled samples in the
support set in classification and use $K=5$ or $K=10$ labeled points with $2$
gradient steps of parameter update in regression. As shown in Figure 7a and
7b, Ensemble and Joint-Training do not always lead to improvements, and the
performance gain is minor. Omni-Training instead outperforms all the compared
methods consistently, which demonstrates that the proposed Omni-Net and the
Omni-Loss designs provide a better solution to bridge pre-training and meta-
training and acquire both domain transferability and task transferability.

Extension to Other Algorithms. Despite the competitive performance on various
benchmarks, we also want to demonstrate that different few-shot learning
algorithms can benefit from the Omni-Training framework. We extend Omni-
Training to more algorithms. Since most pre-training algorithms adopt the
similar pre-training and fine-tuning process, we mainly investigate the
varieties of meta-training algorithms including MatchingNet [13], MAML [9] and
RelationNet [42]. We conduct experiments in the mini-ImageNet dataset since
some algorithms cannot deal with the regression problem. As shown in Figure
7c, Omni-Training with different algorithms significantly outperforms the
corresponding baselines. This demonstrates that our framework can generally
accommodate different few-shot learning algorithms.

(a)  Figure 9: Attention comparison of the representations from the pre-
training model, the meta-training model and the three data flows in Omni-
Training.

Comparison of Each Flow with Baselines. We investigate whether the
coordination of pre-training and meta-training with the shared parameters in
our _tri-flow_ architecture can improve the performance of specific flows.
Figure 8a reports the training losses and validation accuracies of the _pre-
flow_ in Omni-Training and pre-training algorithm Baseline [15] alone, while
Figure 8b reports the results of the _meta-flow_ in Omni-Training and the
meta-training algorithm ProtoNet [14]. The experiments are conducted in the
CUB dataset with $K=5$. The pre-flow and the meta-flow in Omni-Training reach
lower losses and higher accuracies than the baselines trained independently.
Even though the pre-flow and Baseline achieve nearly the same training loss,
the pre-flow achieves much higher validation accuracy than Baseline. This
shows that the knowledge communication enables pre-flow to obtain part of task
transferability and meta-flow to obtain part of domain transferability to
improve their performance.

We also compare the transferability of the pre-training method and the pre-
flow on the mini-ImageNet and mini-ImageNet$\rightarrow$CUB datasets. As shown
in Figure 8c, the pre-flow also outperforms pre-training in various
situations. We further investigate fine-tuning the representations with $1/10$
gradient steps. The performance of the pre-training model drops a lot with
limited updates, but the pre-flow still performs well and comparably with the
pre-training model updated $10$ more times. This reveals that the pre-flow
also acquires task transferability to fast adapt across tasks. These results
demonstrate that Omni-Training coordinates the two parallel flows and makes
each gain the other kind of transferability.

(a) CUB

(b) mini-ImageNet$\rightarrow$CUB

(c) Parameter Sensitivity

Figure 10: (a)(b) Results of modifying different numbers of layers into Omni-Layers on the CUB dataset and the mini-ImageNet$\rightarrow$CUB dataset respectively. (c) The sensitivity of the performance with respect to the hyper-parameters $\lambda$. TABLE V: Ablation study on the losses in the Omni-Training framework. $\mathcal{L}_{\texttt{pre}}$ | $\mathcal{L}_{\texttt{meta}}$ | $\mathcal{L}_{\texttt{joint}}$ | $\mathcal{R}$ | ImageNet | CUB | ImageNet$\rightarrow$CUB
---|---|---|---|---|---|---
✓ | - | - | - | $74.37$ | $82.80$ | $65.57$
- | ✓ | - | - | $73.66$ | $87.91$ | $62.14$
- | - | ✓ | - | $76.95$ | $85.15$ | $66.09$
✓ | ✓ | - | - | $78.12$ | $87.34$ | $69.66$
✓ | ✓ | ✓ | - | $80.79$ | $88.54$ | $70.12$
✓ | ✓ | ✓ | ✓ | $81.26$ | $91.09$ | $71.30$
TABLE VI: Ablation study on the model size. Method | #Params | ImageNet | CUB | ImageNet$\rightarrow$CUB
---|---|---|---|---
ProtoNet | 11.17M | $73.68$ | $87.42$ | $62.02$
ProtoNet* | 13.98M | $73.44$ | $87.81$ | $61.27$
Omni-Training | 13.98M | $81.26$ | $91.09$ | $71.30$

Comparison of Attention Maps. We compare the spatial attention in different
representations learned by pre-training, meta-training and the three data
flows in Omni-Training. From Figure 9, we observe that pre-training
representations focus on a broad area containing the objects as well as some
noisy context, which fully grab the domain knowledge but lack some
concentration on the important information to discriminate different
categories. On the contrary, the meta-training representations focus on a very
small area with very concise information, which is easy to generalize across
tasks quickly but also easy to make mistakes when the attending area deviates
only a little from the objects. Such deviation is more likely to occur with
the domain shift. Such attention heatmaps are consistent with our analyses
before that pre-training learns representations with higher domain
transferability while meta-training learns representations with higher task
transferability.

Switching to Omni-Training, the pre-flow focuses on a more concise area only
including the whole object while ignoring the noisy context. The meta-flow
focuses on a broader area to grab more knowledge in the whole domain and
increase the tolerance of mistakes. This observation demonstrates that there
is knowledge transfer between pre-flow and meta-flow, which coordinates these
two flows and improves them with the other kind of transferability. The joint-
flow shows a different attention map from the pre-flow and the meta-flow. This
also demonstrates that the three flows in the Omni-Training framework focus on
different areas on the input space and form a more comprehensive understanding
of the datapoints.

###  7.2 Framework Analysis

Ablation Study of Losses. We conduct an ablation study by using different
combinations of losses in the Omni-Training framework. For the losses of
$\mathcal{L}_{\texttt{pre}}$, $\mathcal{L}_{\texttt{meta}}$ and
$\mathcal{L}_{\texttt{joint}}$, if we do not use any of the three losses, we
will not use the corresponding branch for inference. We report results on
mini-ImageNet, CUB and mini-ImageNet$\rightarrow$CUB datasets with $K=5$ in
Table V. We observe that all of the loss functions in the tri-flow design
including the self-distillation regularization contribute to the improvement
of the Omni-Training framework.

Influence of the Model Size. In Omni-Net, we use lightweight $1\times 1$
convolution layers for the parallel branches. Although the number of
parameters does not increase significantly (from $11.17$M to $13.98$M if we
use ResNet-18), there is still a concern that the performance gain of Omni-
Training may come from the increase in the model size. Thus, we add the same
parameters as these additional $1\times 1$ convolution layers to the original
ResNet-18 backbone, and denote it as ResNet-18*. Though having the same number
of parameters, ResNet-18* is different from our Omni-Training backbone because
it does not have different data flows inside respectively for pre-training and
meta-training, and is only trained with one learning paradigm. We train
ProtoNet [14] with the ResNet-18* backbone (denoted as ProtoNet*) and report
the accuracy with the support set size $K=5$ in Table VI.

Despite having more parameters, ProtoNet* does not show obvious improvement
over ProtoNet. This indicates that simply increasing the model complexity does
not ensure better performance. Omni-Training has comparable parameters with
ProtoNet*, but outperforms ProtoNet* with a large margin. This reveals that
the main reason that improves the performance is not increasing the model
size, but coordinating pre-training and meta-training to learn deep
representations with both domain transferability and task transferability.

Backbone Modification. We investigate the incluence of the number of Omni-
Layers of the backbone. Since ResNet-18 is composed of $8$ Res-Blocks, we
attempt to keep the first $n$ Res-Blocks unchanged and transform the rest
$8-n$ blocks into Omni-Layers. The first index of the block with Omni-Layers
is $n+1$. We train the models with these modified backbones. We report
classification results with $K=5$ in the CUB dataset (Figure 10a) and the
mini-ImageNet$\rightarrow$CUB dataset (Figure 10b). When the index of the
first block with Omni-Layers is $1$, which means the whole backbone is changed
into Omni-Net, the model performs best. As the index increases, which means
more preceding layers are completely shared between different flows as done in
Multi-Task Learning, the accuracy drops sharply. This reveals the efficacy of
the Omni-Layers on learning the three flows to coordinate pre-training and
meta-training. Omni-Net is a general-purpose backbone for few-shot learning.

Parameter Sensitivity. We analyze the sensitivity of the loss trade-off hyper-
parameter $\lambda$. We report the accuracy on the mini-ImageNet dataset with
$K=1$ and on the cross-domain mini-ImageNet$\rightarrow$CUB dataset with $K=5$
in Figure 10c. We observe that the model performs well in a range of
parameters: $[1.0,3.0]$. However, the performance degrades when setting
$\lambda=0$, _i.e._ , removing the self-distillation regularization. In
general, we use the same hyper-parameter: $\lambda=3.0$ for the different
tasks in our experiments to avoid over-tuning it.

##  8 Conclusion

This paper focuses on learning transferable representations for few-shot
learning, which enables the model to fast generalize to new domains and tasks
with a few examples. We pinpoint that domain transferability and task
transferability are the key factors to data-efficiency in downstream tasks. We
further empirically show that pre-training and meta-training methods and
simple combinations of them cannot obtain both domain transferability and task
transferability, so we propose Omni-Training to bridge pre-training and meta-
training with both types of transferability. With the tri-flow Omni-Net
architecture, the model preserves the specific transferability of pre-training
and meta-training and coordinates these flows by routing their representations
via the joint-flow, making each gain the other kind of transferability. We
design an Omni-Loss to learn the three flows and impose a self-distillation
regularization to enable knowledge transfer across the training process. Omni-
Training is a general framework that accommodates various existing pre-
training and meta-training algorithms. Thorough evaluations on cross-task and
cross-domain datasets in classification, regression and reinforcement learning
problems shows that Omni-Training consistently and clearly outperforms the
state-of-the-art deep learning methods for few-shot learning.

## Acknowledgments

This work was supported by the National Megaproject for New Generation AI
(2020AAA0109201), National Natural Science Foundation of China (62022050 and
62021002), Beijing Nova Program (Z201100006820041), and BNRist Innovation Fund
(BNR2021RC01002).

## References

  * [1] Y. LeCun, Y. Bengio, and G. Hinton, “Deep learning,” _nature_ , vol. 521, no. 7553, pp. 436–444, 2015.
  * [2] K. He, X. Zhang, S. Ren, and J. Sun, “Deep residual learning for image recognition,” in _CVPR_ , 2016, pp. 770–778.
  * [3] D. Silver, A. Huang, C. J. Maddison, A. Guez, L. Sifre, G. Van Den Driessche, J. Schrittwieser, I. Antonoglou, V. Panneershelvam, M. Lanctot _et al._ , “Mastering the game of go with deep neural networks and tree search,” _nature_ , vol. 529, no. 7587, pp. 484–489, 2016.
  * [4] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova, “Bert: Pre-training of deep bidirectional transformers for language understanding,” _arXiv preprint_ , 2018.
  * [5] D. Adiwardana, M.-T. Luong, D. R. So, J. Hall, N. Fiedel, R. Thoppilan, Z. Yang, A. Kulshreshtha, G. Nemade, Y. Lu _et al._ , “Towards a human-like open-domain chatbot,” _arXiv preprint_ , 2020.
  * [6] R. Bommasani, P. Liang _et al._ , “On the opportunities and risks of foundation models,” _arXiv preprint_ , 2021.
  * [7] C. Yu, J. Liu, and S. Nemati, “Reinforcement learning in healthcare: A survey,” _arXiv preprint_ , 2019.
  * [8] L. Fei-Fei, R. Fergus, and P. Perona, “One-shot learning of object categories,” _TPAMI_ , vol. 28, no. 4, pp. 594–611, 2006.
  * [9] C. Finn, P. Abbeel, and S. Levine, “Model-agnostic meta-learning for fast adaptation of deep networks,” in _ICML_ , 2017, pp. 1126–1135.
  * [10] X. Wang, J. Gao, M. Long, and J. Wang, “Self-tuning for data-efficient deep learning,” in _ICML_ , 2021, pp. 10 738–10 748.
  * [11] A. Kolesnikov, L. Beyer, X. Zhai, J. Puigcerver, J. Yung, S. Gelly, and N. Houlsby, “Big transfer (bit): General visual representation learning,” in _ECCV_ , 2020, pp. 491–507.
  * [12] A. Radford, K. Narasimhan, T. Salimans, and I. Sutskever, “Improving language understanding by generative pre-training,” _OpenAI blog_ , 2018.
  * [13] O. Vinyals, C. Blundell, T. Lillicrap, D. Wierstra _et al._ , “Matching networks for one shot learning,” in _NeurIPS_ , 2016, pp. 3630–3638.
  * [14] J. Snell, K. Swersky, and R. Zemel, “Prototypical networks for few-shot learning,” in _NeurIPS_ , 2017, pp. 4077–4087.
  * [15] W.-Y. Chen, Y.-C. Liu, Z. Kira, Y.-C. F. Wang, and J.-B. Huang, “A closer look at few-shot classification,” in _ICLR_ , 2019.
  * [16] Y. Guo, N. C. Codella, L. Karlinsky, J. V. Codella, J. R. Smith, K. Saenko, T. Rosing, and R. Feris, “A broader study of cross-domain few-shot learning,” in _ECCV_ , 2020, pp. 124–141.
  * [17] V. Dumoulin, N. Houlsby, U. Evci, X. Zhai, R. Goroshin, S. Gelly, and H. Larochelle, “Comparing transfer and meta learning approaches on a unified few-shot classification benchmark,” _arXiv preprint_ , 2021.
  * [18] Y. Wang, Q. Yao, J. T. Kwok, and L. M. Ni, “Generalizing from a few examples: A survey on few-shot learning,” _ACM Computing Surveys_ , vol. 53, no. 3, pp. 1–34, 2020.
  * [19] Y. Yu, “Towards sample efficient reinforcement learning.” in _IJCAI_ , 2018, pp. 5739–5743.
  * [20] C. Szegedy, V. Vanhoucke, S. Ioffe, J. Shlens, and Z. Wojna, “Rethinking the inception architecture for computer vision,” in _CVPR_ , 2016, pp. 2818–2826.
  * [21] K. He, R. Girshick, and P. Dollár, “Rethinking imagenet pre-training,” in _ICCV_ , 2019, pp. 4918–4927.
  * [22] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, I. Sutskever _et al._ , “Language models are unsupervised multitask learners,” _OpenAI blog_ , vol. 1, no. 8, p. 9, 2019.
  * [23] T. Brown, B. Mann, N. Ryder, M. Subbiah, J. D. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal, A. Herbert-Voss, G. Krueger, T. Henighan, R. Child, A. Ramesh, D. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei, “Language models are few-shot learners,” in _NeurIPS_ , 2020, pp. 1877–1901.
  * [24] Z. Cao, M. Kwon, and D. Sadigh, “Transfer reinforcement learning across homotopy classes,” _IEEE Robotics and Automation Letters_ , vol. 6, no. 2, pp. 2706–2713, 2021.
  * [25] Z. Zhu, K. Lin, and J. Zhou, “Transfer learning in deep reinforcement learning: A survey,” _arXiv preprint_ , 2020.
  * [26] T. Xie, N. Jiang, H. Wang, C. Xiong, and Y. Bai, “Policy finetuning: Bridging sample-efficient offline and online reinforcement learning,” _arXiv preprint_ , 2021.
  * [27] V. Campos, P. Sprechmann, S. S. Hansen, A. Barreto, S. Kapturowski, A. Vitvitskyi, A. P. Badia, and C. Blundell, “Beyond fine-tuning: Transferring behavior in reinforcement learning,” in _ICML Workshop on Unsupervised Reinforcement Learning_ , 2021.
  * [28] M. Schwarzer, N. Rajkumar, M. Noukhovitch, A. Anand, L. Charlin, D. Hjelm, P. Bachman, and A. Courville, “Pretraining representations for data-efficient reinforcement learning,” _arXiv preprint_ , 2021.
  * [29] N. Dvornik, J. Mairal, and C. Schmid, “Diversity with cooperation: Ensemble methods for few-shot classification,” in _ICCV_ , 2019, pp. 3722–3730.
  * [30] Y. Tian, Y. Wang, D. Krishnan, J. B. Tenenbaum, and P. Isola, “Rethinking few-shot image classification: a good embedding is all you need?” in _ECCV_ , 2020, pp. 266–282.
  * [31] S. Qiao, C. Liu, W. Shen, and A. L. Yuille, “Few-shot image recognition by predicting parameters from activations,” in _CVPR_ , 2018, pp. 7229–7238.
  * [32] H. Qi, M. Brown, and D. G. Lowe, “Low-shot learning with imprinted weights,” in _CVPR_ , 2018, pp. 5822–5830.
  * [33] L. Xuhong, Y. Grandvalet, and F. Davoine, “Explicit inductive bias for transfer learning with convolutional networks,” in _ICML_ , 2018.
  * [34] X. Li, H. Xiong, H. Wang, Y. Rao, L. Liu, Z. Chen, and J. Huan, “Delta: Deep learning transfer using feature map with attention for convolutional networks,” in _ICLR_ , 2019.
  * [35] X. Chen, S. Wang, B. Fu, M. Long, and J. Wang, “Catastrophic forgetting meets negative transfer: Batch spectral shrinkage for safe transfer learning,” in _NeurIPS_ , 2019.
  * [36] K. You, Z. Kou, M. Long, and J. Wang, “Co-tuning for transfer learning,” in _NeurIPS_ , 2020.
  * [37] J. Schmidhuber, “Evolutionary principles in self-referential learning, or on learning how to learn: the meta-meta-… hook,” Ph.D. dissertation, Technische Universität München, 1987.
  * [38] D. K. Naik and R. J. Mammone, “Meta-neural networks that learn by learning,” in _IJCNN_ , 1992, pp. 437–442.
  * [39] S. Thrun and L. Pratt, _Learning to learn_. Springer Science & Business Media, 1998.
  * [40] J. Lu, P. Gong, J. Ye, and C. Zhang, “Learning from very few samples: A survey,” _arXiv preprint_ , 2020.
  * [41] G. Koch, R. Zemel, R. Salakhutdinov _et al._ , “Siamese neural networks for one-shot image recognition,” in _ICML deep learning workshop_ , 2015\.
  * [42] F. Sung, Y. Yang, L. Zhang, T. Xiang, P. H. Torr, and T. M. Hospedales, “Learning to compare: Relation network for few-shot learning,” in _CVPR_ , 2018, pp. 1199–1208.
  * [43] K. Allen, E. Shelhamer, H. Shin, and J. Tenenbaum, “Infinite mixture prototypes for few-shot learning,” in _ICML_ , 2019, pp. 232–241.
  * [44] V. Garcia and J. Bruna, “Few-shot learning with graph neural networks,” in _ICLR_ , 2018.
  * [45] R. Hou, H. Chang, M. Bingpeng, S. Shan, and X. Chen, “Cross attention network for few-shot classification,” in _NeurIPS_ , 2019, pp. 4005–4016.
  * [46] B. Oreshkin, P. R. López, and A. Lacoste, “Tadam: Task dependent adaptive metric for improved few-shot learning,” in _NeurIPS_ , 2018, pp. 721–731.
  * [47] S. W. Yoon, J. Seo, and J. Moon, “Tapnet: Neural network augmented with task-adaptive projection for few-shot learning,” in _ICML_ , 2019, pp. 7115–7123.
  * [48] H.-J. Ye, H. Hu, D.-C. Zhan, and F. Sha, “Few-shot learning via embedding adaptation with set-to-set functions,” in _CVPR_ , 2020.
  * [49] Y. Bengio, S. Bengio, and J. Cloutier, “Learning a synaptic learning rule: Université de montréal,” _Département d’informatique et de recherche opérationnelle_ , 1990.
  * [50] J. Schmidhuber, “Learning to control fast-weight memories: An alternative to dynamic recurrent networks,” _Neural Computation_ , vol. 4, no. 1, pp. 131–139, 1992.
  * [51] M. Andrychowicz, M. Denil, S. Gomez, M. W. Hoffman, D. Pfau, T. Schaul, B. Shillingford, and N. De Freitas, “Learning to learn by gradient descent by gradient descent,” in _NeurIPS_ , 2016, pp. 3981–3989.
  * [52] S. Ravi and H. Larochelle, “Optimization as a model for few-shot learning,” in _ICLR_ , 2017.
  * [53] Z. Li, F. Zhou, F. Chen, and H. Li, “Meta-sgd: Learning to learn quickly for few-shot learning,” _arXiv preprint_ , 2017.
  * [54] Z. Xu, H. van Hasselt, and D. Silver, “Meta-gradient reinforcement learning,” _arXiv preprint_ , 2018.
  * [55] R. Houthooft, R. Y. Chen, P. Isola, B. C. Stadie, F. Wolski, J. Ho, and P. Abbeel, “Evolved policy gradients,” _arXiv preprint_ , 2018.
  * [56] A. Santoro, S. Bartunov, M. Botvinick, D. Wierstra, and T. Lillicrap, “Meta-learning with memory-augmented neural networks,” in _ICML_ , 2016, pp. 1842–1850.
  * [57] T. Munkhdalai and H. Yu, “Meta networks,” in _ICML_ , 2017, pp. 2554–2563.
  * [58] N. Mishra, M. Rohaninejad, X. Chen, and P. Abbeel, “A simple neural attentive meta-learner,” in _ICLR_ , 2018.
  * [59] T. Munkhdalai, X. Yuan, S. Mehri, and A. Trischler, “Rapid adaptation with conditionally shifted neurons,” in _ICML_ , 2018, pp. 3664–3673.
  * [60] Y. Duan, J. Schulman, X. Chen, P. L. Bartlett, I. Sutskever, and P. Abbeel, “Rl2: Fast reinforcement learning via slow reinforcement learning,” _arXiv preprint_ , 2016.
  * [61] J. X. Wang, Z. Kurth-Nelson, D. Tirumala, H. Soyer, J. Z. Leibo, R. Munos, C. Blundell, D. Kumaran, and M. Botvinick, “Learning to reinforcement learn,” _arXiv preprint_ , 2016.
  * [62] Y. Lee and S. Choi, “Gradient-based meta-learning with learned layerwise metric and subspace,” in _ICML_ , 2018, pp. 2927–2936.
  * [63] H. Yao, Y. Wei, J. Huang, and Z. Li, “Hierarchically structured meta-learning,” in _ICML_ , 2019, pp. 7045–7054.
  * [64] Y. Duan, M. Andrychowicz, B. C. Stadie, J. Ho, J. Schneider, I. Sutskever, P. Abbeel, and W. Zaremba, “One-shot imitation learning,” in _NeurIPS_ , 2017.
  * [65] C. Finn, T. Yu, T. Zhang, P. Abbeel, and S. Levine, “One-shot visual imitation learning via meta-learning,” in _CoRL_ , 2017, pp. 357–368.
  * [66] K. Frans, J. Ho, X. Chen, P. Abbeel, and J. Schulman, “Meta learning shared hierarchies,” _arXiv preprint_ , 2017.
  * [67] M. A. Jamal and G.-J. Qi, “Task agnostic meta-learning for few-shot learning,” in _CVPR_ , 2019, pp. 11 719–11 727.
  * [68] A. Xie, A. Singh, S. Levine, and C. Finn, “Few-shot goal inference for visuomotor learning and planning,” in _CoRL_ , 2018, pp. 40–52.
  * [69] A. A. Rusu, D. Rao, J. Sygnowski, O. Vinyals, R. Pascanu, S. Osindero, and R. Hadsell, “Meta-learning with latent embedding optimization,” in _ICLR_ , 2019.
  * [70] L. Bertinetto, J. F. Henriques, P. Torr, and A. Vedaldi, “Meta-learning with differentiable closed-form solvers,” in _ICLR_ , 2019.
  * [71] K. Lee, S. Maji, A. Ravichandran, and S. Soatto, “Meta-learning with differentiable convex optimization,” in _CVPR_ , 2019, pp. 10 657–10 665.
  * [72] H.-Y. Tseng, H.-Y. Lee, J.-B. Huang, and M.-H. Yang, “Cross-domain few-shot classification via learned feature-wise transformation,” in _ICLR_ , 2020\.
  * [73] Q. Sun, Y. Liu, T.-S. Chua, and B. Schiele, “Meta-transfer learning for few-shot learning,” in _CVPR_ , 2019, pp. 403–412.
  * [74] E. Triantafillou, T. Zhu, V. Dumoulin, P. Lamblin, U. Evci, K. Xu, R. Goroshin, C. Gelada, K. Swersky, P. Manzagol, and H. Larochelle, “Meta-dataset: A dataset of datasets for learning to learn from few examples,” in _ICLR_ , 2020.
  * [75] R. S. Sutton, D. A. McAllester, S. P. Singh, and Y. Mansour, “Policy gradient methods for reinforcement learning with function approximation,” in _NeurIPS_ , 2000, pp. 1057–1063.
  * [76] R. S. Sutton and A. G. Barto, _Reinforcement Learning: An Introduction_. The MIT Press, 2018.
  * [77] B. Liu, Y. Cao, Y. Lin, Q. Li, Z. Zhang, M. Long, and H. Hu, “Negative margin matters: Understanding margin in few-shot classification,” in _ECCV_ , 2020, pp. 438–455.
  * [78] A. Afrasiyabi, J.-F. Lalonde, and C. Gagné, “Persistent mixture model networks for few-shot image classification,” _arXiv preprint_ , 2020.
  * [79] H. Wang, H. Zhao, and B. Li, “Bridging multi-task learning and meta-learning: Towards efficient training and effective adaptation,” in _ICML_ , 2021.
  * [80] R. Ni, M. Goldblum, A. Sharaf, K. Kong, and T. Goldstein, “Data augmentation for meta-learning,” in _ICML_ , 2021, pp. 8152–8161.
  * [81] C. Wah, S. Branson, P. Welinder, P. Perona, and S. Belongie, “The caltech-ucsd birds-200-2011 dataset,” 2011.
  * [82] O. Russakovsky, J. Deng, H. Su, J. Krause, S. Satheesh, S. Ma, Z. Huang, A. Karpathy, A. Khosla, M. Bernstein _et al._ , “Imagenet large scale visual recognition challenge,” _IJCV_ , vol. 115, no. 3, pp. 211–252, 2015\.
  * [83] N. Hilliard, L. Phillips, S. Howland, A. Yankov, C. D. Corley, and N. O. Hodas, “Few-shot learning with metric-agnostic conditional embeddings,” _arXiv preprint_ , 2018.
  * [84] J. Krause, M. Stark, J. Deng, and L. Fei-Fei, “3d object representations for fine-grained categorization,” in _ICCV_ , 2013, pp. 554–561.
  * [85] B. Zhou, A. Lapedriza, A. Khosla, A. Oliva, and A. Torralba, “Places: A 10 million image database for scene recognition,” _TPAMI_ , vol. 40, no. 6, pp. 1452–1464, 2017.
  * [86] G. V. Horn, O. M. Aodha, Y. Song, Y. Cui, C. Sun, A. Shepard, H. Adam, P. Perona, and S. J. Belongie, “The inaturalist species classification and detection dataset,” in _CVPR_ , 2018, pp. 8769–8778.
  * [87] D. P. Kingma and J. Ba, “Adam: A method for stochastic optimization,” _arXiv preprint_ , 2014.
  * [88] G. Blanchard, G. Lee, and C. Scott, “Generalizing from several related classification tasks to a new unlabeled sample,” in _NeurIPS_ , 2011, pp. 2178–2186.
  * [89] Y. Duan, X. Chen, R. Houthooft, J. Schulman, and P. Abbeel, “Benchmarking deep reinforcement learning for continuous control,” in _ICML_ , 2016, pp. 1329–1338.
  * [90] E. Todorov, T. Erez, and Y. Tassa, “Mujoco: A physics engine for model-based control,” in _IROS_ , 2012, pp. 5026–5033.
  * [91] R. J. Williams, “Simple statistical gradient-following algorithms for connectionist reinforcement learning,” _Machine learning_ , vol. 8, no. 3, pp. 229–256, 1992.