Original: ST-MAML : A Stochastic-Task based Method for Task-Heterogeneous Meta-Learning
Revised: ST-MAML : A Stochastic-Task based Method for Task-Heterogeneous Meta-Learning

Original: Abstract
Revised: Abstract

Original: Optimization-based meta-learning typically assumes tasks are sampled from a single distribution – an assumption that oversimpliﬁes and limits the diversity of tasks that meta-learning can model.
Revised: Optimization-based meta-learning typically assumes tasks are sampled from a single distribution – an assumption that oversimpliﬁes and limits the diversity of tasks that meta-learning can model.

Original: Handling tasks from multiple distributions is challenging for meta-learning because it adds ambiguity to task identities.
Revised: Handling tasks from multiple distributions is challenging for meta-learning because it adds ambiguity to task identities.

Original: This paper proposes a novel method, ST-MAML , that empowers modelagnostic meta-learning ( MAML ) to learn from multiple task distributions.
Revised: This paper proposes a novel method, ST-MAML , that empowers modelagnostic meta-learning ( MAML ) to learn from multiple task distributions.

Original: ST-MAML encodes tasks using a stochastic neural network module, that summarizes every task with a stochastic representation.
Revised: ST-MAML encodes tasks using a stochastic neural network module, that summarizes every task with a stochastic representation.

Original: The proposed Stochastic Task ( ST ) strategy learns a distribution of solutions for an ambiguous task and allows a meta-model to self-adapt to the current task.
Revised: The proposed Stochastic Task ( ST ) strategy learns a distribution of solutions for an ambiguous task and allows a meta-model to self-adapt to the current task.

Original: ST-MAML also propagates the task representation to enhance input variable encodings.
Revised: ST-MAML also propagates the task representation to enhance input variable encodings.

Original: Empirically, we demonstrate that ST-MAML outperforms the state-of-the-art on two few-shot image classiﬁcation tasks, one curve regression benchmark, one image completion problem, and a real-world temperature prediction application.
Revised: Empirically, we demonstrate that ST-MAML outperforms the state-of-the-art on two few-shot image classiﬁcation tasks, one curve regression benchmark, one image completion problem, and a real-world temperature prediction application.

Original:  INTRODUCTION
Revised:  INTRODUCTION

Original: Meta-learning aims to train a model on multiple machine learning tasks to adapt to a new task with only a few training samples.
Revised: Meta-learning aims to train a model on multiple machine learning tasks to adapt to a new task with only a few training samples.

Original: Optimization-based meta-learning like model-agnostic meta-learning (MAML) facilitate such a goal by involving the optimization process.
Revised: Optimization-based meta-learning like model-agnostic meta-learning (MAML) facilitate such a goal by involving the optimization process.

Original: For example,
Revised: For example,

Original: MAML trains a global initialization of model parameters that are close to the optimal parameter values of every task (Finn et al., 2017).
Revised: MAML trains a global initialization of model parameters that are close to the optimal parameter values of every task (Finn et al., 2017).

Original: Recent methods expand MAML’s “global initialization” to a notion of“globally shared knowledge”, including not only initialization (Finn et al., 2017;
Revised: Recent methods expand MAML’s “global initialization” to a notion of“globally shared knowledge”, including not only initialization (Finn et al., 2017;

Original: Li et al., 2017; Rajeswaran et al., 2019) but also update rules (Andrychowicz et al., 2016; Ravi and Larochelle,
Revised: Li et al., 2017; Rajeswaran et al., 2019) but also update rules (Andrychowicz et al., 2016; Ravi and Larochelle,

Original: [Figure]
Revised: [Figure]

Original: Globally shared knowledge allows these methods to produce good generalization performance on new tasks with a small number of training samples.
Revised: Globally shared knowledge allows these methods to produce good generalization performance on new tasks with a small number of training samples.

Original: Most optimization-based meta-learning algorithms assume all tasks T are identically and independently sampled from a single distribution (Andrychowicz et al., 2016; Finn et al., 2017; Li et al., 2017; Ravi and Larochelle, 2017; Rusu et al., 2018).
Revised: Most optimization-based meta-learning algorithms assume all tasks T are identically and independently sampled from a single distribution (Andrychowicz et al., 2016; Finn et al., 2017; Li et al., 2017; Ravi and Larochelle, 2017; Rusu et al., 2018).

Original: We refer to meta-learning’s task distribution as the “meta-distribution”.
Revised: We refer to meta-learning’s task distribution as the “meta-distribution”.

Original: Formally, these methods assume T ∼
Revised: Formally, these methods assume T ∼

Original: P ( T ) .
Revised: P ( T ) .

Original: Real-world tasks, however, may come from multiple meta-distributions, T ∼ { P 1 ( T ) , P 2 ( T ) , · · · , P k ( T ) } .
Revised: Real-world tasks, however, may come from multiple meta-distributions, T ∼ { P 1 ( T ) , P 2 ( T ) , · · · , P k ( T ) } .

Original: For instance, when analyzing multiple writers’ hand written digits, writers from different age group (like children versus adults) indicate different meta-distributions.
Revised: For instance, when analyzing multiple writers’ hand written digits, writers from different age group (like children versus adults) indicate different meta-distributions.

Original: This more challenging setup, we call task heterogeneity, poses technical challenges to homogenous strategies like MAML (Vuorio et al., 2019).
Revised: This more challenging setup, we call task heterogeneity, poses technical challenges to homogenous strategies like MAML (Vuorio et al., 2019).

Original: For task heterogeneity, a naive and widely accepted metalearning solution ﬁrst learns a globally shared initialization across all meta-distributions and then tailors model parameters to the current task (Vuorio et al., 2019; Yao et al.,
Revised: For task heterogeneity, a naive and widely accepted metalearning solution ﬁrst learns a globally shared initialization across all meta-distributions and then tailors model parameters to the current task (Vuorio et al., 2019; Yao et al.,

Original: 2020, 2019; Lee and Choi, 2018; Oreshkin et al., 2018).
Revised: 2020, 2019; Lee and Choi, 2018; Oreshkin et al., 2018).

Original: The tailoring step needs to rely on the task-speciﬁc information or, ideally, the identity information of the task.
Revised: The tailoring step needs to rely on the task-speciﬁc information or, ideally, the identity information of the task.

Original: It, therefore, requires the meta-learner to infer the potential identity of a new task from a limited number of annotated samples (Finn et al., 2018).
Revised: It, therefore, requires the meta-learner to infer the potential identity of a new task from a limited number of annotated samples (Finn et al., 2018).

Original: This requirement raises severe uncertainty issues – a challenge known as “task ambiguity”.
Revised: This requirement raises severe uncertainty issues – a challenge known as “task ambiguity”.

Original: Figure 1 provides a concrete example of the task ambiguity that arises from limited annotated data and unknown meta distribution when facing task heterogeneity.
Revised: Figure 1 provides a concrete example of the task ambiguity that arises from limited annotated data and unknown meta distribution when facing task heterogeneity.

Original: Surprisingly, recent optimization-based meta-learning literature pays little attention to the task ambiguity challenge (Vuorio et al., 2019; Yao et al., 2020, 2019; Lee and Choi, 2018).
Revised: Surprisingly, recent optimization-based meta-learning literature pays little attention to the task ambiguity challenge (Vuorio et al., 2019; Yao et al., 2020, 2019; Lee and Choi, 2018).

Original: Besides, the task heterogeneity ampliﬁes the “distribution shift” issue (Zhang et al., 2021; Dubey et al., 2021).
Revised: Besides, the task heterogeneity ampliﬁes the “distribution shift” issue (Zhang et al., 2021; Dubey et al., 2021).

Original: The difference between two tasks can signiﬁcantly increase in the heterogeneous setup since tasks are from various meta-distributions.
Revised: The difference between two tasks can signiﬁcantly increase in the heterogeneous setup since tasks are from various meta-distributions.

Original: This paper proposes a novel meta-learning method
Revised: This paper proposes a novel meta-learning method

Original: ST-MAML for the task heterogeneity challenge.
Revised: ST-MAML for the task heterogeneity challenge.

Original: Our approach extends MAML by modeling tasks as a stochastic variable that we call the “stochastic task”.
Revised: Our approach extends MAML by modeling tasks as a stochastic variable that we call the “stochastic task”.

Original: Stochastic tasks (STs) let us learn a distribution of solutions to capture the uncertainty of an ambiguous new task.
Revised: Stochastic tasks (STs) let us learn a distribution of solutions to capture the uncertainty of an ambiguous new task.

Original: At the same time,
Revised: At the same time,

Original: STs enable self-adaptive model initialization based on the current task.
Revised: STs enable self-adaptive model initialization based on the current task.

Original: We use variational inference as a solver and the whole learning process is meta-distribution agnostic.
Revised: We use variational inference as a solver and the whole learning process is meta-distribution agnostic.

Original: We apply ST-MAML to a wide range of common meta-learning benchmarks including synthetic regression, image completion, and few-shot image classiﬁcation, where ST-MAML exceeds the performance of existing work.
Revised: We apply ST-MAML to a wide range of common meta-learning benchmarks including synthetic regression, image completion, and few-shot image classiﬁcation, where ST-MAML exceeds the performance of existing work.

Original: We also build a large temperature prediction dataset that highlights the challenges of real-world meta-distributions.
Revised: We also build a large temperature prediction dataset that highlights the challenges of real-world meta-distributions.

Original: Our empirical results demonstrate that ST-MAML outperforms the MAML baselines by 40% on this new task.
Revised: Our empirical results demonstrate that ST-MAML outperforms the MAML baselines by 40% on this new task.

Original:  METHODS 2.1 PRELIMINARIES ON META LEARNING
Revised:  METHODS 2.1 PRELIMINARIES ON META LEARNING

Original: We describe a supervised learning task in meta-learning as
Revised: We describe a supervised learning task in meta-learning as

Original: T =
Revised: T =

Original: {L
Revised: {L

Original: [Equation]
Revised: [Equation]

Original: f
Revised: f

Original: [Equation]
Revised: [Equation]

Original: T
Revised: T

Original: [Equation]
Revised: [Equation]

Original: D
Revised: D

Original: [Equation]
Revised: [Equation]

Original: Here L oss () , which takes as input model f θ T and dataset, describes the loss function that measures the quality of learner f θ T , whose parameter weight is θ T .
Revised: Here L oss () , which takes as input model f θ T and dataset, describes the loss function that measures the quality of learner f θ T , whose parameter weight is θ T .

Original: Every task includes an annotated training set D tr T =
Revised: Every task includes an annotated training set D tr T =

Original: [ X tr T , Y tr T ] and a test set D te T =
Revised: [ X tr T , Y tr T ] and a test set D te T =

Original: [ X te T , Y te T ] .
Revised: [ X te T , Y te T ] .

Original: During meta-training, the test set D te T is fully observed, but during meta-testing only its input X te T is available.
Revised: During meta-training, the test set D te T is fully observed, but during meta-testing only its input X te T is available.

Original: D tr T and D te T are sampled from X ×Y ,
Revised: D tr T and D te T are sampled from X ×Y ,

Original: X describes the input space and Y is the output space.
Revised: X describes the input space and Y is the output space.

Original: The goal of meta learning is to train a learning machine which can perform well on D te T after ﬁne-tuning on this task’s training set D tr T .
Revised: The goal of meta learning is to train a learning machine which can perform well on D te T after ﬁne-tuning on this task’s training set D tr T .

Original: The difﬁculty lies at ﬁnding a balance between underﬁtting to all tasks and overﬁtting to any particular task.
Revised: The difﬁculty lies at ﬁnding a balance between underﬁtting to all tasks and overﬁtting to any particular task.

Original: MAML (Finn et al., 2017) achieves such a goal by learning a globally shared weight initialization θ ∗ that is close to the optimal weight parameter of every task.
Revised: MAML (Finn et al., 2017) achieves such a goal by learning a globally shared weight initialization θ ∗ that is close to the optimal weight parameter of every task.

Original: We can write its training objective for getting the best initialization θ ∗ as:
Revised: We can write its training objective for getting the best initialization θ ∗ as:

Original: [Equation]
Revised: [Equation]

Original: where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: MAML samples a set of tasks {T } from the meta distribution P ( T ) and initialize each task’s weight θ 0 T from the global knowledge θ (to be learnt): i.e., setting θ 0 T = θ .
Revised: MAML samples a set of tasks {T } from the meta distribution P ( T ) and initialize each task’s weight θ 0 T from the global knowledge θ (to be learnt): i.e., setting θ 0 T = θ .

Original: On each task, the learner performs gradient descent on its training set
Revised: On each task, the learner performs gradient descent on its training set

Original: D tr T to reach task-speciﬁc ﬁne-tuned parameters θ 1 T .
Revised: D tr T to reach task-speciﬁc ﬁne-tuned parameters θ 1 T .

Original: The test set D te T of task T is used for evaluating parameter θ 1 T , and the evaluation will be used as the objective to optimize for learning the best global knowledge θ .
Revised: The test set D te T of task T is used for evaluating parameter θ 1 T , and the evaluation will be used as the objective to optimize for learning the best global knowledge θ .

Original: In probabilistic language, the above objective (in Eq.
Revised: In probabilistic language, the above objective (in Eq.

Original: (2)) can be equivalently framed as maximizing the likelihood: max
Revised: (2)) can be equivalently framed as maximizing the likelihood: max

Original: [Equation]
Revised: [Equation]

Original: where p ( θ 1 T | D tr T , θ ) is a Dirac distribution derived by minimizing the negative log-likelihood(NLL) on D tr T with gradient descent.
Revised: where p ( θ 1 T | D tr T , θ ) is a Dirac distribution derived by minimizing the negative log-likelihood(NLL) on D tr T with gradient descent.

Original: 2.2 PREVIOUS HETEROGENEOUS META LEARNING
Revised: 2.2 PREVIOUS HETEROGENEOUS META LEARNING

Original: Task-homogeneous meta-learning assumes that there exists one meta-distribution P ( T ) and all tasks are identically and independently (i.i.d.) sampled from P ( T ) .
Revised: Task-homogeneous meta-learning assumes that there exists one meta-distribution P ( T ) and all tasks are identically and independently (i.i.d.) sampled from P ( T ) .

Original: Differently, in a task-heterogeneous setup, there exist multiple metadistributions T ∼ { P 1 ( T ) , P 2 ( T ) , · · · , P k ( T ) } .
Revised: Differently, in a task-heterogeneous setup, there exist multiple metadistributions T ∼ { P 1 ( T ) , P 2 ( T ) , · · · , P k ( T ) } .

Original: We can naively use MAML and assign all tasks with the same global initialization (though they come from different distributions).
Revised: We can naively use MAML and assign all tasks with the same global initialization (though they come from different distributions).

Original: Figure 1(a, b) show that the “task ambiguity” issue is more critical in task-heterogeneous setup and will hinder the generalization from MAML initialization since multiple very different task distributions exist.
Revised: Figure 1(a, b) show that the “task ambiguity” issue is more critical in task-heterogeneous setup and will hinder the generalization from MAML initialization since multiple very different task distributions exist.

Original: A handful of previous works learn a customized initialization that was tailored from global initialization, in order to tackle the task heterogeneity challenge.
Revised: A handful of previous works learn a customized initialization that was tailored from global initialization, in order to tackle the task heterogeneity challenge.

Original: MMAML (Vuorio et al., 2019) learns a deterministic task embedding with an
Revised: MMAML (Vuorio et al., 2019) learns a deterministic task embedding with an

Original: RNN module.
Revised: RNN module.

Original: HSML (Yao et al., 2019) manually designs a task clustering algorithm to assign tasks to different clusters, then customizes the global initialization to each cluster.
Revised: HSML (Yao et al., 2019) manually designs a task clustering algorithm to assign tasks to different clusters, then customizes the global initialization to each cluster.

Original: ARML (Yao et al., 2020) models global knowledge and task-speciﬁc knowledge as graphs; the interaction between tasks is modeled by message passing.
Revised: ARML (Yao et al., 2020) models global knowledge and task-speciﬁc knowledge as graphs; the interaction between tasks is modeled by message passing.

Original: However, none of the recent works consider the task ambiguity issue when solving task-heterogeneous domains.
Revised: However, none of the recent works consider the task ambiguity issue when solving task-heterogeneous domains.

Original: Most frameworks are still based on the assumption that only one path to the posterior in meta-training path to the prior in meta-testing ﬁne-tuning only observed in meta-training observation Θ 0𝒯 =
Revised: Most frameworks are still based on the assumption that only one path to the posterior in meta-training path to the prior in meta-testing ﬁne-tuning only observed in meta-training observation Θ 0𝒯 =

Original: [ θ 0𝒯 , h 0𝒯 ] θ θ 0𝒯 h 0𝒯
Revised: [ θ 0𝒯 , h 0𝒯 ] θ θ 0𝒯 h 0𝒯

Original: X te 𝒯 Z 𝒯
Revised: X te 𝒯 Z 𝒯

Original: X tr 𝒯 Y tr 𝒯 θ 1𝒯 h 1𝒯
Revised: X tr 𝒯 Y tr 𝒯 θ 1𝒯 h 1𝒯

Original: Y te 𝒯 ( , (
Revised: Y te 𝒯 ( , (

Original: X te 𝒯 Y te 𝒯 ( , ( r ′  𝒯 r 𝒯 Θ 1𝒯 =
Revised: X te 𝒯 Y te 𝒯 ( , ( r ′  𝒯 r 𝒯 Θ 1𝒯 =

Original: [ θ 1𝒯 , h 1𝒯 ]
Revised: [ θ 1𝒯 , h 1𝒯 ]

Original: 1: Input: Meta-distribution set { P 1 ( T ) , ··· ,P k ( T ) } , Hyper-parameters γ 1 and γ 2 . 2: Randomly
Revised: 1: Input: Meta-distribution set { P 1 ( T ) , ··· ,P k ( T ) } , Hyper-parameters γ 1 and γ 2 . 2: Randomly

Original: [Table]
Revised: [Table]

Original: parameter
Revised: parameter

Original: [Table] [Table]
Revised: [Table] [Table]

Original: φ ,
Revised: φ ,

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table]
Revised: [Table]

Original: not DONE do 4:
Revised: not DONE do 4:

Original: Sample
Revised: Sample

Original: [Table] [Table]
Revised: [Table] [Table]

Original: do
Revised: do

Original: [Table]
Revised: [Table]

Original: Infer
Revised: Infer

Original: [Table]
Revised: [Table]

Original: of
Revised: of

Original: [Table] [Table]
Revised: [Table] [Table]

Original: |
Revised: |

Original: [Table]
Revised: [Table]

Original: ) and sample z
Revised: ) and sample z

Original: [Table]
Revised: [Table]

Original: ∼ q ( Z T |T ) .
Revised: ∼ q ( Z T |T ) .

Original: [Table]
Revised: [Table]

Original: and
Revised: and

Original: [Table]
Revised: [Table]

Original: Tailor θ with sample z T to get
Revised: Tailor θ with sample z T to get

Original: [Table]
Revised: [Table]

Original: -speciﬁc initialization θ 0 T .
Revised: -speciﬁc initialization θ 0 T .

Original: [eq.(12)] 8:
Revised: [eq.(12)] 8:

Original: [Table]
Revised: [Table]

Original: the raw input.
Revised: the raw input.

Original: [eq.(13)] 9:
Revised: [eq.(13)] 9:

Original: Evaluate the inner loss
Revised: Evaluate the inner loss

Original: [Table]
Revised: [Table]

Original: in ( T ) on training set D tr T .
Revised: in ( T ) on training set D tr T .

Original: [eq.(17)] 10:
Revised: [eq.(17)] 10:

Original: Compute
Revised: Compute

Original: [Table]
Revised: [Table]

Original: and
Revised: and

Original: [Table]
Revised: [Table]

Original: with
Revised: with

Original: [Table]
Revised: [Table]

Original: descent [eq.(18)]: θ
Revised: descent [eq.(18)]: θ

Original: [Table]
Revised: [Table]

Original: =
Revised: =

Original: [Table]
Revised: [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: ) , h
Revised: ) , h

Original: [Table]
Revised: [Table]

Original: =
Revised: =

Original: [Table] [Table]
Revised: [Table] [Table]

Original: 11: end for
Revised: 11: end for

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: , w , β ] (cid:80) T L ELBO ( T ) .
Revised: , w , β ] (cid:80) T L ELBO ( T ) .

Original: [eq.(16)] 13: end while distribution exists to explain a task’s observed training set (e.g., a new task should be assigned to only one cluster in
Revised: [eq.(16)] 13: end while distribution exists to explain a task’s observed training set (e.g., a new task should be assigned to only one cluster in

Original: HSML).
Revised: HSML).

Original: However, the source of a task can be highly uncertain based on limited annotated data.
Revised: However, the source of a task can be highly uncertain based on limited annotated data.

Original: Figure 1(b) shows that there can be multiple explanations of an observed dataset in the task-heterogeneous setup and we should not expect to obtain a unique solution.
Revised: Figure 1(b) shows that there can be multiple explanations of an observed dataset in the task-heterogeneous setup and we should not expect to obtain a unique solution.

Original: 2.3 STOCHASTIC Z T TO ENCODE TASK
Revised: 2.3 STOCHASTIC Z T TO ENCODE TASK

Original: When facing the task-heterogeneous setup, we hypothesize that a meta-learner that can encode potential tasks’ patterns will help alleviate the task ambiguity issue.
Revised: When facing the task-heterogeneous setup, we hypothesize that a meta-learner that can encode potential tasks’ patterns will help alleviate the task ambiguity issue.

Original: These patterns could describe valuable information about tasks like the more possible shapes of curves for a regression meta-application.
Revised: These patterns could describe valuable information about tasks like the more possible shapes of curves for a regression meta-application.

Original: Moreover, we propose to enable task encoding with uncertainty estimates.
Revised: Moreover, we propose to enable task encoding with uncertainty estimates.

Original: This is because learning a task representation from its limited annotated data is challenging and such uncertainty measures can help inform the downstream meta-adaptation to new tasks (see Figure 1(b)).
Revised: This is because learning a task representation from its limited annotated data is challenging and such uncertainty measures can help inform the downstream meta-adaptation to new tasks (see Figure 1(b)).

Original: This hypothesis motivates us to describe a task T with a stochastic variable Z T and model its distribution to condition on observations.
Revised: This hypothesis motivates us to describe a task T with a stochastic variable Z T and model its distribution to condition on observations.

Original: With this additional latent variable, we can rewrite the per task likelihood L ( T ) in Eq.
Revised: With this additional latent variable, we can rewrite the per task likelihood L ( T ) in Eq.

Original: (3) as:
Revised: (3) as:

Original: L
Revised: L

Original: [Equation]
Revised: [Equation]

Original: (cid:88)
Revised: (cid:88)

Original: [Equation]
Revised: [Equation]

Original: We assume in the second term from above, Z T only conditions on D tr T .
Revised: We assume in the second term from above, Z T only conditions on D tr T .

Original: Figure 2 shows our design.
Revised: Figure 2 shows our design.

Original: In Section (2.5), we show that the likelihood is intractable as deﬁned above, and choose to maximize its evidence lower bound (a.k.a ELBO) instead.
Revised: In Section (2.5), we show that the likelihood is intractable as deﬁned above, and choose to maximize its evidence lower bound (a.k.a ELBO) instead.

Original: Optimizing this variational objective requires the prior p ( Z T | D tr T ) and the posterior q ( Z T |T ) .
Revised: Optimizing this variational objective requires the prior p ( Z T | D tr T ) and the posterior q ( Z T |T ) .

Original: We model the prior p ( Z T | D tr T ) as a Gaussian distribution, whose mean and variance are outputs from a multi-layer perceptron (MLP) module with input vector r T :
Revised: We model the prior p ( Z T | D tr T ) as a Gaussian distribution, whose mean and variance are outputs from a multi-layer perceptron (MLP) module with input vector r T :

Original: [Equation]
Revised: [Equation]

Original: Here vector r T is a vector summarizing the encoding of a task T .
Revised: Here vector r T is a vector summarizing the encoding of a task T .

Original: We propose a neural network module to learn r T from the sample observations D tr T .
Revised: We propose a neural network module to learn r T from the sample observations D tr T .

Original: The training observations of task T consist of unordered annotated data pairs [( x tr T , y tr T )] .
Revised: The training observations of task T consist of unordered annotated data pairs [( x tr T , y tr T )] .

Original: Permutation invariance is a desirable property for functions acting on sets.
Revised: Permutation invariance is a desirable property for functions acting on sets.

Original: Zaheer et al.
Revised: Zaheer et al.

Original: (2017) showed any function acting on sets S is permutation invariant if and only if it can be decomposed as ρ ( (cid:80) s ∈ S φ ( s )) for suitable choice of transformations ρ, φ .
Revised: (2017) showed any function acting on sets S is permutation invariant if and only if it can be decomposed as ρ ( (cid:80) s ∈ S φ ( s )) for suitable choice of transformations ρ, φ .

Original: We follow such a design, and encode a task by encoding every pair of its observation in D tr T through a neural network layer:
Revised: We follow such a design, and encode a task by encoding every pair of its observation in D tr T through a neural network layer:

Original: [Equation]
Revised: [Equation]

Original: Eq.
Revised: Eq.

Original: (8) uses average function as aggregation operator to obtain the task embedding because it is able to remove the inductive bias due to different sizes of training set from r T .
Revised: (8) uses average function as aggregation operator to obtain the task embedding because it is able to remove the inductive bias due to different sizes of training set from r T .

Original: In Eq.
Revised: In Eq.

Original: (7), g Enc φ () is implemented as a MLP module with learnable parameter φ .
Revised: (7), g Enc φ () is implemented as a MLP module with learnable parameter φ .

Original: We then approximate the intractable posterior distribution q ( Z T |T ) of Z T as conditioned on the whole { D tr T , D te T } (see Section (2)): [Equation]
Revised: We then approximate the intractable posterior distribution q ( Z T |T ) of Z T as conditioned on the whole { D tr T , D te T } (see Appendix S2):

Original: [Equation]
Revised: [Equation]

Original: where |T | = | D tr T | + | D te T | , µ ( · ) and σ ( · ) are the same
Revised: where |T | = | D tr T | + | D te T | , µ ( · ) and σ ( · ) are the same

Original: MLP modules we have in Eq.
Revised: MLP modules we have in Eq.

Original: 2.4 ST-MAML : SELF ADAPTATION WITH Z T
Revised: 2.4 ST-MAML : SELF ADAPTATION WITH Z T

Original: We propose to revise MAML for the heterogeneous metalearning setup using the summary task representation Z T , creating ST-MAML .
Revised: We propose to revise MAML for the heterogeneous metalearning setup using the summary task representation Z T , creating ST-MAML .

Original: Z T helps tailor the global initialization θ to task-speciﬁc initialization θ 0 T for a task T .
Revised: Z T helps tailor the global initialization θ to task-speciﬁc initialization θ 0 T for a task T .

Original: Its basic motivation is to improves ﬂexibility by incorporating task information into the model.
Revised: Its basic motivation is to improves ﬂexibility by incorporating task information into the model.

Original: This self adaption design is motivated by the recent ideas that design self-adaptation conditioned on global knowledge to conquer distribution shift issue in domain generalization/adaptation (Zhang et al., 2021;
Revised: This self adaption design is motivated by the recent ideas that design self-adaptation conditioned on global knowledge to conquer distribution shift issue in domain generalization/adaptation (Zhang et al., 2021;

Original: Dubey et al., 2021; Xiao et al., 2021; Vuorio et al., 2019).
Revised: Dubey et al., 2021; Xiao et al., 2021; Vuorio et al., 2019).

Original: Thereexist many potential ways to use Z T to tailor the global initialization θ to task-speciﬁc initialization θ 0 T .
Revised: Thereexist many potential ways to use Z T to tailor the global initialization θ to task-speciﬁc initialization θ 0 T .

Original: We assume our target learning machine is a composition of a base learner and a task learner:
Revised: We assume our target learning machine is a composition of a base learner and a task learner:

Original: [Equation]
Revised: [Equation]

Original: Here the base learner’s parameters are θ b , and the task learner’s parameters are θ c .
Revised: Here the base learner’s parameters are θ b , and the task learner’s parameters are θ c .

Original: For example, in an image classiﬁcation domain the base learner would be the the CNN backbone and the task learner would be the last linear layer.
Revised: For example, in an image classiﬁcation domain the base learner would be the the CNN backbone and the task learner would be the last linear layer.

Original: We can then rewrite θ =
Revised: We can then rewrite θ =

Original: [ θ b , θ c ] .
Revised: [ θ b , θ c ] .

Original: We propose to only customize θ c with Z T :
Revised: We propose to only customize θ c with Z T :

Original: [Equation]
Revised: [Equation]

Original: Here z T is sampled from the distribution q ( Z T |T ) during meta-training and from p ( Z T | D tr T ) during meta-testing.
Revised: Here z T is sampled from the distribution q ( Z T |T ) during meta-training and from p ( Z T | D tr T ) during meta-testing.

Original: σ is the sigmoid function, (cid:12) represents the element-wise multiplication, w =
Revised: σ is the sigmoid function, (cid:12) represents the element-wise multiplication, w =

Original: [ w 1 , w 0 ] T are learnable parameters.
Revised: [ w 1 , w 0 ] T are learnable parameters.

Original: w gate , the gate vector will apply element-wise scaling to navigate global initialization θ to task-speciﬁc initialization θ 0 T .
Revised: w gate , the gate vector will apply element-wise scaling to navigate global initialization θ to task-speciﬁc initialization θ 0 T .

Original: Moreover, we design additional customized knowledge for task T .
Revised: Moreover, we design additional customized knowledge for task T .

Original: The basic intuition is that the ﬁnal prediction of a meta-learner depends on both model parameters and input representations.
Revised: The basic intuition is that the ﬁnal prediction of a meta-learner depends on both model parameters and input representations.

Original: To increase the capacity of the task-speciﬁc knowledge, we propose to further propagate task representation Z T into encoding augmented feature representations we denote as h T .
Revised: To increase the capacity of the task-speciﬁc knowledge, we propose to further propagate task representation Z T into encoding augmented feature representations we denote as h T .

Original: We concatenate h T with a sample’s input representation x T , and feed the combined vector ˆ x T to our learning machine as its new input.
Revised: We concatenate h T with a sample’s input representation x T , and feed the combined vector ˆ x T to our learning machine as its new input.

Original: [Equation]
Revised: [Equation]

Original: Same as Eq.
Revised: Same as Eq.

Original: (12), z T is sampled from its distribution, β =
Revised: (12), z T is sampled from its distribution, β =

Original: [ β 1 , β 0 ] are learnable parameters.
Revised: [ β 1 , β 0 ] are learnable parameters.

Original: Now when facing a new task T , a meta-model will ﬁrst generate the task-speciﬁc knowledge that includes both augmented feature h T and task-speciﬁc parameter θ T .
Revised: Now when facing a new task T , a meta-model will ﬁrst generate the task-speciﬁc knowledge that includes both augmented feature h T and task-speciﬁc parameter θ T .

Original: We denote the combined knowledge set for task T as:
Revised: We denote the combined knowledge set for task T as:

Original: [Equation]
Revised: [Equation]

Original: This is the meta-knowledge we need to learn in ST-MAML .
Revised: This is the meta-knowledge we need to learn in ST-MAML .

Original: We note its initial values as Θ 0 T =
Revised: We note its initial values as Θ 0 T =

Original: [ θ 0 T , h 0 T ] and ﬁne-tuned values as Θ 1 T =
Revised: [ θ 0 T , h 0 T ] and ﬁne-tuned values as Θ 1 T =

Original: [ θ 1 T , h 1 T ] .
Revised: [ θ 1 T , h 1 T ] .

Original: Aiming to learn the meta knowledge deﬁned in Eq.
Revised: Aiming to learn the meta knowledge deﬁned in Eq.

Original: (14), we can rewrite our objective (task likelihood) in Eq.
Revised: (14), we can rewrite our objective (task likelihood) in Eq.

Original: [Equation]
Revised: [Equation]

Original: This follows the Bayesian graph provided in Figure 2.
Revised: This follows the Bayesian graph provided in Figure 2.

Original: Design Choices: There exist many other possible probabilistic designs besides Figure 2.
Revised: Design Choices: There exist many other possible probabilistic designs besides Figure 2.

Original: For instance, we can model every variable in the ﬁgure as a stochastic distribution and build a complicated hybrid framework.
Revised: For instance, we can model every variable in the ﬁgure as a stochastic distribution and build a complicated hybrid framework.

Original: However, this will lead to excessive stochasticity and increase the potential of underﬁtting in a limited data situation.
Revised: However, this will lead to excessive stochasticity and increase the potential of underﬁtting in a limited data situation.

Original: Instead, similar to p ( Θ 1 T | Θ 0 T , D tr T ) , we choose to model both p ( h 0 T | Z T ) and p ( θ 0 T | θ , Z T ) as deterministic (see Eq.
Revised: Instead, similar to p ( Θ 1 T | Θ 0 T , D tr T ) , we choose to model both p ( h 0 T | Z T ) and p ( θ 0 T | θ , Z T ) as deterministic (see Eq.

Original: (12) and Eq.
Revised: (12) and Eq.

Original: that allow us to employ an amortized variational inference technique Ravi and Beatson (2019).
Revised: that allow us to employ an amortized variational inference technique Ravi and Beatson (2019).

Original: Our design is different from recent probabilistic extensions of MAML Finn et al.
Revised: Our design is different from recent probabilistic extensions of MAML Finn et al.

Original: ; Yoon et al.
Revised: ; Yoon et al.

Original: They conduct inference on model parameters θ T (initial value θ 0 T or ﬁne-tuned value θ 1 T ).
Revised: They conduct inference on model parameters θ T (initial value θ 0 T or ﬁne-tuned value θ 1 T ).

Original: Our ST-MAML shifts the burden of variational inference to the task representation Z T , whose dimension is of multiple orders smaller than the size of model parameters.
Revised: Our ST-MAML shifts the burden of variational inference to the task representation Z T , whose dimension is of multiple orders smaller than the size of model parameters.

Original: 2.5 ST-MAML : UPDATE RULES
Revised: 2.5 ST-MAML : UPDATE RULES

Original: Variational Objective: To optimize the intractable likelihood as deﬁned in Eq.
Revised: Variational Objective: To optimize the intractable likelihood as deﬁned in Eq.

Original: (15), we choose to maximize its evidence lower bound (a.k.a ELBO) instead:
Revised: (15), we choose to maximize its evidence lower bound (a.k.a ELBO) instead:

Original: [Equation]
Revised: [Equation]

Original: During meta-training, we sample m tasks and optimize the empirical average 1 m m (cid:80)
Revised: During meta-training, we sample m tasks and optimize the empirical average 1 m m (cid:80)

Original: [Equation]
Revised: [Equation]

Original: Update Rules: Much like MAML, the optimization of
Revised: Update Rules: Much like MAML, the optimization of

Original: ST-MAML consists of two nested loops.
Revised: ST-MAML consists of two nested loops.

Original: Figure 3 shows the iterative optimization process.
Revised: Figure 3 shows the iterative optimization process.

Original: In the inner loop, for the j th training data, we concatenate x tr T ,j with augmented feature h 0 T to get augmented input vector ˆ x tr T ,j .
Revised: In the inner loop, for the j th training data, we concatenate x tr T ,j with augmented feature h 0 T to get augmented input vector ˆ x tr T ,j .

Original: We feed ˆ x tr T ,j into the learning machine f
Revised: We feed ˆ x tr T ,j into the learning machine f

Original: whose parameter is θ 0 T to calculate the inner loss:
Revised: whose parameter is θ 0 T to calculate the inner loss:

Original: [Equation]
Revised: [Equation]

Original: The inner loss is then used for updating θ 0 T and h 0 T :
Revised: The inner loss is then used for updating θ 0 T and h 0 T :

Original: [Equation]
Revised: [Equation]

Original: Figure 3 shows we can optimize the inner loss for K iterations to achieve a closer approximation for optimal values in Eq.
Revised: Figure 3 shows we can optimize the inner loss for K iterations to achieve a closer approximation for optimal values in Eq.

Original: In the outer loop, we maximize the approximated ELBO L ELBO in Eq.
Revised: In the outer loop, we maximize the approximated ELBO L ELBO in Eq.

Original: (16) using a batch of m
Revised: (16) using a batch of m

Original: 𝒯 ( D tr 𝒯 , D te 𝒯 )
Revised: 𝒯 ( D tr 𝒯 , D te 𝒯 )

Original: Z 𝒯 ( C , θ ) ( h 0𝒯 , X tr 𝒯 ) Cat
Revised: Z 𝒯 ( C , θ ) ( h 0𝒯 , X tr 𝒯 ) Cat

Original: ̂ X
Revised: ̂ X

Original: tr 𝒯 θ 0𝒯 ℒ in ( f θ 0 𝒯 , [ ̂ X tr 𝒯 , Y tr 𝒯 ])
Revised: tr 𝒯 θ 0𝒯 ℒ in ( f θ 0 𝒯 , [ ̂ X tr 𝒯 , Y tr 𝒯 ])

Original: Tailor ( h 1𝒯 , X tr 𝒯 ) Cat ̂ X tr 𝒯 θ 1𝒯 ( h K 𝒯 , X te 𝒯 )
Revised: Tailor ( h 1𝒯 , X tr 𝒯 ) Cat ̂ X tr 𝒯 θ 1𝒯 ( h K 𝒯 , X te 𝒯 )

Original: Cat θ K 𝒯 ̂ X te 𝒯 ⋯ ⋯ KL ( q ( Z 𝒯 | 𝒯)
Revised: [Figure]

Original: || p ( Z 𝒯 | D tr 𝒯 )) ℒ ELBO (
Revised: θ

Original: [Figure]
Revised: [Figure]

Revised: 𝒯 ̂

Original: [Figure]
Revised: [Figure]

Original: in ( f
Revised: ⋯ ⋯

Original: [Figure]
Revised: [Figure]

Original: , [ ̂ X tr
Revised: [Figure]

Original: [Figure]
Revised: [Figure]

Original: , Y tr 𝒯 ]) ℒ
Revised: D tr 𝒯 ))

Original: [Figure]
Revised: [Figure]

Original: tasks.
Revised: tasks.

Original: The amortized variational technique allows us to conduct the sampling from q ( Θ 1 T |T ) by ﬁrst sampling from q ( Z T |T ) and then applying a deterministic transformation using Eq.
Revised: The amortized variational technique allows us to conduct the sampling from q ( Θ 1 T |T ) by ﬁrst sampling from q ( Z T |T ) and then applying a deterministic transformation using Eq.

Original: (12) and Eq.
Revised: (12) and Eq.

Original: Algorithm of ST-MAML : We describe the procedure of
Revised: Algorithm of ST-MAML : We describe the procedure of

Original: ST-MAML in the form of pseudo code as shown in Algorithm 1.
Revised: ST-MAML in the form of pseudo code as shown in Algorithm 1.

Original: Note, parameters of neural functions µ ( · ) , σ ( · ) , g Enc φ () , g Gate w () , and g
Revised: Note, parameters of neural functions µ ( · ) , σ ( · ) , g Enc φ () , g Gate w () , and g

Original: In β () are updated in the outer loop.
Revised: In β () are updated in the outer loop.

Original: Theoretical Analysis of ST-MAML :
Revised: Theoretical Analysis of ST-MAML :

Original: We also provide a second interpretation of our objective from an information bottleneck perspective and prove they lead to exactly the same target.
Revised: We also provide a second interpretation of our objective from an information bottleneck perspective and prove they lead to exactly the same target.

Original: See Appendix 3 for detailed proofs.
Revised: See Appendix S3 for detailed proofs.

Original: 2.6 CONNECTING TO RELATED WORK
Revised: 2.6 CONNECTING TO RELATED WORK

Original: Optimization-based meta-learning methods facilitate the model’s adaption to new tasks through global knowledge learned by the optimization process.
Revised: Optimization-based meta-learning methods facilitate the model’s adaption to new tasks through global knowledge learned by the optimization process.

Original: Meta-LSTM (Ravi and Larochelle, 2017) meta-learns the update rule with an RNN meta-learner.
Revised: Meta-LSTM (Ravi and Larochelle, 2017) meta-learns the update rule with an RNN meta-learner.

Original: MAML (Finn et al., 2017) trains a global initialization close to the optimal value of every task.
Revised: MAML (Finn et al., 2017) trains a global initialization close to the optimal value of every task.

Original: Leveraging diverse meta-knowledge further accelerates the learning process.
Revised: Leveraging diverse meta-knowledge further accelerates the learning process.

Original: In Meta-SGD (Li et al., 2017), the meta-knowledge consists of both initialization and learning rate.
Revised: In Meta-SGD (Li et al., 2017), the meta-knowledge consists of both initialization and learning rate.

Original: ALFA (Baik et al., 2020) proposes to meta-learn both initialization and hyperparameter update module.
Revised: ALFA (Baik et al., 2020) proposes to meta-learn both initialization and hyperparameter update module.

Original: Most methods assign the same global knowledge to every task that leads to sub-optimal solutions for heterogeneous settings.
Revised: Most methods assign the same global knowledge to every task that leads to sub-optimal solutions for heterogeneous settings.

Original: Besides, they are all deterministic and can only learn one solution for a new task.
Revised: Besides, they are all deterministic and can only learn one solution for a new task.

Original: Bayesian approaches are a long-standing discipline that incorporates uncertainty in modeling.
Revised: Bayesian approaches are a long-standing discipline that incorporates uncertainty in modeling.

Original: Multiple recent works extend MAML into the Bayesian framework and recast meta-learning as the probabilistic framework (Finn et al., 2018; Grant et al., 2018;
Revised: Multiple recent works extend MAML into the Bayesian framework and recast meta-learning as the probabilistic framework (Finn et al., 2018; Grant et al., 2018;

Original: Yoon et al., 2018; Ravi and Beatson, 2019; Garnelo et al., 2018b).
Revised: Yoon et al., 2018; Ravi and Beatson, 2019; Garnelo et al., 2018b).

Original: BMAML (Yoon et al., 2018) recast MAML into probabilistic framework and provides a Bayesian explanation of MAML.
Revised: BMAML (Yoon et al., 2018) recast MAML into probabilistic framework and provides a Bayesian explanation of MAML.

Original: PLATIPUS (Finn et al., 2018) builds upon amortized variational inference and injects Gaussian noise into the gradient during the metatesting time to learn a distribution over model parameters.
Revised: PLATIPUS (Finn et al., 2018) builds upon amortized variational inference and injects Gaussian noise into the gradient during the metatesting time to learn a distribution over model parameters.

Original: LLAMA (Grant et al., 2018) applies Laplace approximation for modeling the parameter distribution, but it requires the approximation of a high dimensional covariance matrix.
Revised: LLAMA (Grant et al., 2018) applies Laplace approximation for modeling the parameter distribution, but it requires the approximation of a high dimensional covariance matrix.

Original: These methods view model parameters (i.e. network weights and bias) as random variables and perform inference on them.
Revised: These methods view model parameters (i.e. network weights and bias) as random variables and perform inference on them.

Original: This leads to signiﬁcant challenges when working with complicated models and high-dimensional data.
Revised: This leads to signiﬁcant challenges when working with complicated models and high-dimensional data.

Original: Our work also loosely connects to the “prototype metalearning" (Triantaﬁllou et al., 2019;
Revised: Our work also loosely connects to the “prototype metalearning" (Triantaﬁllou et al., 2019;

Original: Snell et al., 2017).
Revised: Snell et al., 2017).

Original: These studies learn a prototype for every class we need to predict and the ﬁnal prediction depends on the distances between instances and prototypes.
Revised: These studies learn a prototype for every class we need to predict and the ﬁnal prediction depends on the distances between instances and prototypes.

Original: Amortized bayesian prototype meta-learning (Sun et al., 2021) assumes a distribution over class prototypes.
Revised: Amortized bayesian prototype meta-learning (Sun et al., 2021) assumes a distribution over class prototypes.

Original: This design requires prior knowledge about the classes of tasks and only applies to the classiﬁcation homogeneous-meta setup.
Revised: This design requires prior knowledge about the classes of tasks and only applies to the classiﬁcation homogeneous-meta setup.

Original: Another line of related works studies neural approximators of the stochastic process family (Garnelo et al., 2018b;
Revised: Another line of related works studies neural approximators of the stochastic process family (Garnelo et al., 2018b;

Original: Wang and Van Hoof, 2020; Louizos et al., 2019; Kim et al., 2018).
Revised: Wang and Van Hoof, 2020; Louizos et al., 2019; Kim et al., 2018).

Original: They learn a prior for every task or further use a hierarchical model that learns the instance prior.
Revised: They learn a prior for every task or further use a hierarchical model that learns the instance prior.

Original: However, these methods don’t share knowledge across tasks.
Revised: However, these methods don’t share knowledge across tasks.

Original: In Appendix 1,

Original: Table 1 compares related lines of works with ours.
Revised: Table S compares related lines of works with ours.

Original: [Table]
Revised: [Table]

Original:  EXPERIMENTS
Revised:  EXPERIMENTS

Original: Our experiments are designed to answer the following:
Revised: Our experiments are designed to answer the following:

Original: Q1.
Revised: Q1.

Original: Does ST-MAML successfully meta-learn from heterogeneous tasks across a variety of applications?
Revised: Does ST-MAML successfully meta-learn from heterogeneous tasks across a variety of applications?

Original: Q2.
Revised: Q2.

Original: How does ST-MAML perform when we have more or less task ambiguity?
Revised: How does ST-MAML perform when we have more or less task ambiguity?

Original: Q3.
Revised: Q3.

Original: How does ST-MAML compare to previous heterogeneous meta approaches in terms of accuracy and adaptation?
Revised: How does ST-MAML compare to previous heterogeneous meta approaches in terms of accuracy and adaptation?

Original: Q4.
Revised: Q4.

Original: How does ST-MAML perform when applied to a challenging real-world dataset?
Revised: How does ST-MAML perform when applied to a challenging real-world dataset?

Original: To answer Q1 , we select a wide range of applications in our experiments.
Revised: To answer Q1 , we select a wide range of applications in our experiments.

Original: We provide a summary of our experimental datasets, and their properties in Table 1 .
Revised: We provide a summary of our experimental datasets, and their properties in Table 1 .

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: We compare against several baselines representing four meta-learning groups: (1) meta-learning methods designed for homogeneous tasks: MAML (Finn et al., 2017) and
Revised: We compare against several baselines representing four meta-learning groups: (1) meta-learning methods designed for homogeneous tasks: MAML (Finn et al., 2017) and

Original: MetaSGD (Li et al., 2017).
Revised: MetaSGD (Li et al., 2017).

Original: (2) Meta-learning methods designed for heterogeneous tasks including MMAML (Vuorio et al., 2019) and HSMAML (Yao et al., 2019).
Revised: (2) Meta-learning methods designed for heterogeneous tasks including MMAML (Vuorio et al., 2019) and HSMAML (Yao et al., 2019).

Original: Bayesian meta-learning methods: Bayesian MAML (Yoon et al., 2018), which recasts MAML into the Bayesian framework.
Revised: Bayesian meta-learning methods: Bayesian MAML (Yoon et al., 2018), which recasts MAML into the Bayesian framework.

Original: (4) Neural processes (NPs) methods (Garnelo et al., 2018a,b).
Revised: (4) Neural processes (NPs) methods (Garnelo et al., 2018a,b).

Original: NPs learn a distribution over solutions and are regarded as state-of-the-art methods for small scale metalearning regression applications.
Revised: NPs learn a distribution over solutions and are regarded as state-of-the-art methods for small scale metalearning regression applications.

Original: 3.1 2D REGRESSION: SIMULATED STUDIES
Revised: 3.1 2D REGRESSION: SIMULATED STUDIES

Original: To answer Q2 , we generate synthetic heterogeneous regression tasks that come from multiple functional families of curves.
Revised: To answer Q2 , we generate synthetic heterogeneous regression tasks that come from multiple functional families of curves.

Original: We use probabilistic meta-learning models to sample and visualize multiple solutions.
Revised: We use probabilistic meta-learning models to sample and visualize multiple solutions.

Original: Setup.
Revised: Setup.

Original: We follow a similar setup as (Yao et al., 2020) to generate 2D regression tasks.
Revised: We follow a similar setup as (Yao et al., 2020) to generate 2D regression tasks.

Original: The meta-distribution set { P k ( T ) } consists of 6 function families including sinusoids, straight line, quadratic, cubic, quadratic surface , and ripple functions.
Revised: The meta-distribution set { P k ( T ) } consists of 6 function families including sinusoids, straight line, quadratic, cubic, quadratic surface , and ripple functions.

Original: We perturb the output by adding Gaussian noise with standard deviation 0.3 in meta-train tasks.
Revised: We perturb the output by adding Gaussian noise with standard deviation 0.3 in meta-train tasks.

Original: During metatraining, every task is uniformly randomly sampled from one of the 6 function families, and the size of the training set | D tr T | = 10 .
Revised: During metatraining, every task is uniformly randomly sampled from one of the 6 function families, and the size of the training set | D tr T | = 10 .

Original: We adopted mean square error (MSE) to measure prediction accuracy.
Revised: We adopted mean square error (MSE) to measure prediction accuracy.

Original: A detailed description of the setup and model architecture is available in Appendix 5.
Revised: A detailed description of the setup and model architecture is available in Appendix S5.

Original: Results, ablations, and analysis.
Revised: Results, ablations, and analysis.

Original: We train models on around 10 , 000 tasks and evaluate it on over 1 , 000 newly sampled tasks.
Revised: We train models on around 10 , 000 tasks and evaluate it on over 1 , 000 newly sampled tasks.

Original: The results are summarized in Table 2.
Revised: The results are summarized in Table 2.

Original: We can see clearly ST-MAML outperforms the baselines.
Revised: We can see clearly ST-MAML outperforms the baselines.

Original: To better investigate the contribution of each component, we perform ablation experiments by either removing model tailoring or input variable augmentation.
Revised: To better investigate the contribution of each component, we perform ablation experiments by either removing model tailoring or input variable augmentation.

Original: Table 2 shows that both types of task-speciﬁc knowledge provide important contributions and their combination gives the best performance.
Revised: Table 2 shows that both types of task-speciﬁc knowledge provide important contributions and their combination gives the best performance.

Original: We visualize example curve ﬁts in Figure 4 and Figure(Appendix).
Revised: We visualize example curve ﬁts in Figure 4 and Figure S1.

Original: During meta-testing, we can decrease the size of training set or increase the noise level such that tasks ambiguity can be more concerning.
Revised: During meta-testing, we can decrease the size of training set or increase the noise level such that tasks ambiguity can be more concerning.

Original: In Figure 1(Appendix), all sampled solutions are close to the ground-truth since tasks are less uncertain.
Revised: In Figure S1, all sampled solutions are close to the ground-truth since tasks are less uncertain.

Original: Differently, Figure 4 shows that as tasks become more ambiguous, those sampled solutions by ST-MAML tend to span a wider range.
Revised: Differently, Figure 4 shows that as tasks become more ambiguous, those sampled solutions by ST-MAML tend to span a wider range.

Original: However, they stay faithful around those annotated training data.
Revised: However, they stay faithful around those annotated training data.

Original: More analysis visualization results can be found in Appendix 5.
Revised: More analysis visualization results can be found in Appendix S5.

Original: 3.2 HETEROGENEOUS FEW-SHOT CLASSIFICATION
Revised: 3.2 HETEROGENEOUS FEW-SHOT CLASSIFICATION

Original: To answer Q3 , we apply ST-MAML on two common fewshot classiﬁcation benchmarks from the literature.
Revised: To answer Q3 , we apply ST-MAML on two common fewshot classiﬁcation benchmarks from the literature.

Original: With space limit, results on CelebA data are in Appendix 4.
Revised: With space limit, results on CelebA data are in Appendix S4.

Original: Setup N-way K-shot classiﬁcation is a popular setup in few-shot meta-learning (Chen et al., 2019; Ren et al., 2018;
Revised: Setup N-way K-shot classiﬁcation is a popular setup in few-shot meta-learning (Chen et al., 2019; Ren et al., 2018;

Original: Vinyals et al., 2016).
Revised: Vinyals et al., 2016).

Original: The training set of every task consists of N classes with K labeled data in each class.
Revised: The training set of every task consists of N classes with K labeled data in each class.

Original: In the benchmark Plain-Multi dataset, each meta-task is sampled from one of four diverse datasets (Yao et al., 2019).
Revised: In the benchmark Plain-Multi dataset, each meta-task is sampled from one of four diverse datasets (Yao et al., 2019).

Original: We follow the benchmark architecture, including a feature learner using four convolutional blocks.
Revised: We follow the benchmark architecture, including a feature learner using four convolutional blocks.

Original: Our ST module takes the input x into two convolutional blocks with 6 channels, and then concatenate the output vector with the target variable into a two-layer MLP to model the mean and variance of Z T .
Revised: Our ST module takes the input x into two convolutional blocks with 6 channels, and then concatenate the output vector with the target variable into a two-layer MLP to model the mean and variance of Z T .

Original: Results, ablations, and analysis.
Revised: Results, ablations, and analysis.

Original: After training on over 50 , 000 total tasks, the model is evaluated on 1 , 000 tasks for each dataset and the results are summarized in Table 3.
Revised: After training on over 50 , 000 total tasks, the model is evaluated on 1 , 000 tasks for each dataset and the results are summarized in Table 3.

Original: The most relevant method is MMAML.
Revised: The most relevant method is MMAML.

Original: It learns a deterministic task embedding with an RNN module and encodes all parameters in both base learner f θ b and task learner
Revised: It learns a deterministic task embedding with an RNN module and encodes all parameters in both base learner f θ b and task learner

Original: f θ c .
Revised: f θ c .

Original: Our method outperforms it on every dataset.
Revised: Our method outperforms it on every dataset.

Original: Also, the probabilistic framework enables us to achieve consistently low variance.
Revised: Also, the probabilistic framework enables us to achieve consistently low variance.

Original: Note that HSMAML uses the prior knowledge about the number of clusters, which plays an important role
Revised: Note that HSMAML uses the prior knowledge about the number of clusters, which plays an important role

Original: [Table]
Revised: [Table]

Original: with respect to the ﬁnal accuracy.
Revised: with respect to the ﬁnal accuracy.

Original: Our ST-MAML does not reply upon such prior and achieves lower variance and similar performance than HSMAML.
Revised: Our ST-MAML does not reply upon such prior and achieves lower variance and similar performance than HSMAML.

Original: We again run two ablated versions of the proposed ST-MAML
Revised: We again run two ablated versions of the proposed ST-MAML

Original: [Table]
Revised: [Table]

Original: and compare it against the full version.
Revised: and compare it against the full version.

Original: The combination of input augmentation and model tailoring yields the best results and is most capable of confronting task-heterogeneity
Revised: The combination of input augmentation and model tailoring yields the best results and is most capable of confronting task-heterogeneity

Original: [Table]
Revised: [Table]

Original: 3.3 REAL-WORLD TEMPERATURE PREDICTION
Revised: 3.3 REAL-WORLD TEMPERATURE PREDICTION

Original: Now we answer Q4 by applying ST-MAML to a challenging regression problem using real-world data.
Revised: Now we answer Q4 by applying ST-MAML to a challenging regression problem using real-world data.

Original: Setup.
Revised: Setup.

Original: The NOAA Global Surface Summary of the Day (GSOD) dataset contains daily weather data from thousands of stations around the world.
Revised: The NOAA Global Surface Summary of the Day (GSOD) dataset contains daily weather data from thousands of stations around the world.

Original: Each task is created by sampling data points from (station, year) pairs.
Revised: Each task is created by sampling data points from (station, year) pairs.

Original: Each sample takes in one date of the year along with 15 weather features such as wind speed, station elevation, precipitation, fog, air pressure, etc for that date.
Revised: Each sample takes in one date of the year along with 15 weather features such as wind speed, station elevation, precipitation, fog, air pressure, etc for that date.

Original: It then learns to predict the average temperature in Fahrenheit on that day.
Revised: It then learns to predict the average temperature in Fahrenheit on that day.

Original: We remove important information like the weather station number, name, latitude, and longitude.
Revised: We remove important information like the weather station number, name, latitude, and longitude.

Original: Hiding the station information in this way creates a highly heterogeneous problem where each station generates its own task distribution.
Revised: Hiding the station information in this way creates a highly heterogeneous problem where each station generates its own task distribution.

Original: The model sees 10 days of labeled temperature data before predicting the temperature on 100 test days.
Revised: The model sees 10 days of labeled temperature data before predicting the temperature on 100 test days.

Original: More technical details can be found in Appendix 5.
Revised: More technical details can be found in Appendix S5.

Original: Results and analysis.
Revised: Results and analysis.

Original: After 100 epochs of training on approximately 42 , 000 unique (station, year) tasks, we evaluate the model on a test set of 1 , 000 (station, year) pairs.
Revised: After 100 epochs of training on approximately 42 , 000 unique (station, year) tasks, we evaluate the model on a test set of 1 , 000 (station, year) pairs.

Original: The results are summarized in Table 4.
Revised: The results are summarized in Table 4.

Original: ST-MAML predictions are approximately 40% more accurate than MAML.
Revised: ST-MAML predictions are approximately 40% more accurate than MAML.

Original: MetaSGD, designed for homogeneous meta-learning, achieves low accuracy because the globally learned learning knowledge hurts the model’s generalization to unseen tasks from different distributions.
Revised: MetaSGD, designed for homogeneous meta-learning, achieves low accuracy because the globally learned learning knowledge hurts the model’s generalization to unseen tasks from different distributions.

Original: This is consistent with our assumption that incorporating task-speciﬁc knowledge into the model can help solve the task-heterogeneous challenge.
Revised: This is consistent with our assumption that incorporating task-speciﬁc knowledge into the model can help solve the task-heterogeneous challenge.

Original: We also perform ablation experiments in Table 4.
Revised: We also perform ablation experiments in Table 4.

Original: Both tailored initialization and augmented features outperform the baselines, and they combine for further improvement.
Revised: Both tailored initialization and augmented features outperform the baselines, and they combine for further improvement.

Original: Figure 2 (Appendix) provides a visualization of trained ST-MAML on the NOAA-GSOD temperature prediction task.
Revised: Figure S2 provides a visualization of trained ST-MAML on the NOAAGSOD temperature prediction task.

Original: 3.4 HETEROGENEOUS IMAGE COMPLETION
Revised: 3.4 HETEROGENEOUS IMAGE COMPLETION

Original: While we have already demonstrated ST-MAML on regression and classiﬁcation tasks, we continue to answer Q1 by expanding to image completion, which is a popular small scale meta-learning task.
Revised: While we have already demonstrated ST-MAML on regression and classiﬁcation tasks, we continue to answer Q1 by expanding to image completion, which is a popular small scale meta-learning task.

Original: Setup.
Revised: Setup.

Original: In our heterogeneous image completion application, the meta distribution set { P k ( T ) } = { MNSIT , FMNIST , KMNIST } .
Revised: In our heterogeneous image completion application, the meta distribution set { P k ( T ) } = { MNSIT , FMNIST , KMNIST } .

Original: Every task contains one image of size 28 × 28 sampled randomly from one of the three dataset distributions.
Revised: Every task contains one image of size 28 × 28 sampled randomly from one of the three dataset distributions.

Original: In meta-training, 40 pixels are observed for every image, thus, | D tr T | = 40 .
Revised: In meta-training, 40 pixels are observed for every image, thus, | D tr T | = 40 .

Original: We use coordinates as inputs and pixel value as the target variable.
Revised: We use coordinates as inputs and pixel value as the target variable.

Original: Each image completion can be interpreted as a meta-learning task which generalizes the knowledge from a limited training set | D tr T | = 40 to the entire image of size
Revised: Each image completion can be interpreted as a meta-learning task which generalizes the knowledge from a limited training set | D tr T | = 40 to the entire image of size

Original: | D te T | = 784 .
Revised: | D te T | = 784 .

Original: Architecture details can be found in
Revised: Architecture details can be found in

Original: Appendix 5.
Revised: Appendix S5.

Original: Baselines, results and analysis.
Revised: Baselines, results and analysis.

Original: The described setup is a benchmark task for Neural processes (Garnelo et al., 2018a,b).
Revised: The described setup is a benchmark task for Neural processes (Garnelo et al., 2018a,b).

Original: Thus, we compare our proposed ST-MAML with neural processes (NP) (Garnelo et al., 2018b) and conditional neural processes (CNP) (Garnelo et al., 2018a) which is a deterministic NP.
Revised: Thus, we compare our proposed ST-MAML with neural processes (NP) (Garnelo et al., 2018b) and conditional neural processes (CNP) (Garnelo et al., 2018a) which is a deterministic NP.

Original: The numerical comparison is shown in Table 5.
Revised: The numerical comparison is shown in Table 5.

Original: ST-MAML achieves higher completion precision compared with NP and CNP.
Revised: ST-MAML achieves higher completion precision compared with NP and CNP.

Original: We leave out the variance for all methods because the difference is insigniﬁcant.
Revised: We leave out the variance for all methods because the difference is insigniﬁcant.

Original: Image completion task can be highly ambiguous, because there exist multiple full image choices that could explain the pattern of a handful of observed pixels, especially for gray images.
Revised: Image completion task can be highly ambiguous, because there exist multiple full image choices that could explain the pattern of a handful of observed pixels, especially for gray images.

Original: Uncertainty arises on three levels: the inter-class level, inter-distribution level, and cross-distribution level.
Revised: Uncertainty arises on three levels: the inter-class level, inter-distribution level, and cross-distribution level.

Original: ST-MAML can capture more potential truths by learning a distribution of possibilities rather than a unique mapping.
Revised: ST-MAML can capture more potential truths by learning a distribution of possibilities rather than a unique mapping.

Original: We visualize observations and their completions in Figure 5.
Revised: We visualize observations and their completions in Figure 5.

Original: Interestingly, when we compare the two half-rows describing image completion for a button-up shirt image, the half-row with more pixels observed during meta-testing,
Revised: Interestingly, when we compare the two half-rows describing image completion for a button-up shirt image, the half-row with more pixels observed during meta-testing,

Original: [Table]
Revised: [Table]

Original: [Figure]
Revised: [Figure]

Original: its task is less ambiguous.
Revised: its task is less ambiguous.

Original: Therefore, its completed images are closer to the original image.
Revised: Therefore, its completed images are closer to the original image.

Original: This reﬂects one merit of
Revised: This reﬂects one merit of

Original: ST-MAML : its set operations allows ST-MAML to learn from any size of the training set during meta-testing.
Revised: ST-MAML : its set operations allows ST-MAML to learn from any size of the training set during meta-testing.

Original: [Table]
Revised: [Table]

Original: 3.5 VISUALIZATION OF GATE VECTORS w gate .
Revised: 3.5 VISUALIZATION OF GATE VECTORS w gate .

Original: As noted in Section (2.4) gate vector w gate (Eq 12), which originates from stochastic task variable Z T , translates global initialization θ to task-speciﬁc initialization θ 0 T .
Revised: As noted in Section (2.4) gate vector w gate (Eq 12), which originates from stochastic task variable Z T , translates global initialization θ to task-speciﬁc initialization θ 0 T .

Original: Thus, we hypothesis patterns of gate vectors contain information about the relationships between similar tasks.
Revised: Thus, we hypothesis patterns of gate vectors contain information about the relationships between similar tasks.

Original: To gain insights into the tasks’ gate vectors w gate , we visualize sampled vectors on two applications: 2D regression and image completion.
Revised: To gain insights into the tasks’ gate vectors w gate , we visualize sampled vectors on two applications: 2D regression and image completion.

Original: For both applications, we sample 200 tasks from each P k ( T ) , and visualize their gate vectors w gate using a t-SNE plot (Van der Maaten and Hinton, 2008).
Revised: For both applications, we sample 200 tasks from each P k ( T ) , and visualize their gate vectors w gate using a t-SNE plot (Van der Maaten and Hinton, 2008).

Original: The visualizations are shown in Figure 6.
Revised: The visualizations are shown in Figure 6.

Original: Gate vectors of tasks from the same distribution (shown as same color points) are clustered on t-SNE plots while tasks from very distinct distributions are further away.
Revised: Gate vectors of tasks from the same distribution (shown as same color points) are clustered on t-SNE plots while tasks from very distinct distributions are further away.

Original: For instance, In Figure 6 left, sinusoidal regression tasks (blue) sit far away from ripple surface tasks (brown).
Revised: For instance, In Figure 6 left, sinusoidal regression tasks (blue) sit far away from ripple surface tasks (brown).

Original: These observations can be seen as evidence of the task identiﬁcation capability of ST-MAML .
Revised: These observations can be seen as evidence of the task identiﬁcation capability of ST-MAML .

Original: Furthermore, tasks from similar distributions may entangle (Figure 6 left, linear, quadratic, and cubic regression tasks).
Revised: Furthermore, tasks from similar distributions may entangle (Figure 6 left, linear, quadratic, and cubic regression tasks).

Original: The uncertain identity of similar tasks justiﬁes the representation of task information as stochastic variables.
Revised: The uncertain identity of similar tasks justiﬁes the representation of task information as stochastic variables.

Original:  CONCLUSION
Revised:  CONCLUSION

Original: Task heterogeneity is one critical challenge in meta-learning.
Revised: Task heterogeneity is one critical challenge in meta-learning.

Original: Most meta-learning methods assign the same initialization to every task and fail to handle task heterogeneity.
Revised: Most meta-learning methods assign the same initialization to every task and fail to handle task heterogeneity.

Original: ST-MAML encodes tasks using a stochastic task module with set-based operations for permutation-invariance.
Revised: ST-MAML encodes tasks using a stochastic task module with set-based operations for permutation-invariance.

Original: The probabilistic framework allows us to learn a distribution of solutions for ambiguous tasks and recover better potential task identities.
Revised: The probabilistic framework allows us to learn a distribution of solutions for ambiguous tasks and recover better potential task identities.

Original: This stochastic task design allows for customizing global knowledge with a learned stochastic task distribution.
Revised: This stochastic task design allows for customizing global knowledge with a learned stochastic task distribution.

Original: Empirically, we design extensive experiments on various applications and show that ST-MAML provides an effective way to learn from diverse and ambiguous tasks.
Revised: Empirically, we design extensive experiments on various applications and show that ST-MAML provides an effective way to learn from diverse and ambiguous tasks.

Original: As next step, we plan to add domain generalization during meta-testing to enhance our work.
Revised: As next step, we plan to add domain generalization during meta-testing to enhance our work.