Original: C ONTRASTIVE M ETA -L EARNING FOR P ARTIALLY O BSERVABLE F EW -S HOT L EARNING
Revised: C ONTRASTIVE M ETA -L EARNING FOR P ARTIALLY O BSERVABLE F EW -S HOT L EARNING

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Many contrastive and meta-learning approaches learn representations by identifying common features in multiple views.
Revised: Many contrastive and meta-learning approaches learn representations by identifying common features in multiple views.

Original: However, the formalism for these approaches generally assumes features to be shared across views to be captured coherently.
Revised: However, the formalism for these approaches generally assumes features to be shared across views to be captured coherently.

Original: We consider the problem of learning a unified representation from partial observations , where useful features may be present in only some of the views.
Revised: We consider the problem of learning a unified representation from partial observations , where useful features may be present in only some of the views.

Original: We approach this through a probabilistic formalism enabling views to map to representations with different levels of uncertainty in different components; these views can then be integrated with one another through marginalisation over that uncertainty.
Revised: We approach this through a probabilistic formalism enabling views to map to representations with different levels of uncertainty in different components; these views can then be integrated with one another through marginalisation over that uncertainty.

Original: Our approach, Partial Observation Experts Modelling (POEM), then enables us to meta-learn consistent representations from partial observations.
Revised: Our approach, Partial Observation Experts Modelling (POEM), then enables us to meta-learn consistent representations from partial observations.

Original: We evaluate our approach on an adaptation of a comprehensive few-shot learning benchmark, Meta-Dataset, and demonstrate the benefits of POEM over other meta-learning methods at representation learning from partial observations.
Revised: We evaluate our approach on an adaptation of a comprehensive few-shot learning benchmark, Meta-Dataset, and demonstrate the benefits of POEM over other meta-learning methods at representation learning from partial observations.

Original: We further demonstrate the utility of POEM by meta-learning to represent an environment from partial views observed by an agent exploring the environment.
Revised: We further demonstrate the utility of POEM by meta-learning to represent an environment from partial views observed by an agent exploring the environment.

Original: [Figure]
Revised: [Figure]

Original: This encourages the learning of features that are consistent in all views; in the above example this corresponds to the pattern on the bird’s wing.
Revised: This encourages the learning of features that are consistent in all views; in the above example this corresponds to the pattern on the bird’s wing.

Original: To better handle partial observability, where features may be disjoint between views, we propose Partial Observation Experts Modelling (POEM).
Revised: To better handle partial observability, where features may be disjoint between views, we propose Partial Observation Experts Modelling (POEM).

Original: POEM instead maximises consistency between multiple views, by utilising representation uncertainty to learn which features of the entity are captured by a view, and then combining these representations together by weighting features by their uncertainty via a product of experts model (Hinton, 2002).
Revised: POEM instead maximises consistency between multiple views, by utilising representation uncertainty to learn which features of the entity are captured by a view, and then combining these representations together by weighting features by their uncertainty via a product of experts model (Hinton, 2002).

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Modern contrastive learning methods (Radford et al., 2021; Chen et al., 2020; He et al., 2020; Oord et al., 2019), and embedding-based meta-learning methods such as Prototypical Networks (Snell et al., 2017; Vinyals et al., 2016; Sung et al., 2018; Edwards & Storkey, 2017), learn representations by minimizing a relative distance between representations of related items compared with unrelated items (Ericsson et al., 2021).
Revised: Modern contrastive learning methods (Radford et al., 2021; Chen et al., 2020; He et al., 2020; Oord et al., 2019), and embedding-based meta-learning methods such as Prototypical Networks (Snell et al., 2017; Vinyals et al., 2016; Sung et al., 2018; Edwards & Storkey, 2017), learn representations by minimizing a relative distance between representations of related items compared with unrelated items (Ericsson et al., 2021).

Original: However, we argue that these approaches may learn to disregard potentially relevant features from views that only inform part of the representation in order to achieve better representational consistency, as demonstrated in Figure 1.
Revised: However, we argue that these approaches may learn to disregard potentially relevant features from views that only inform part of the representation in order to achieve better representational consistency, as demonstrated in Figure 1.

Original: We refer to such partially informative views as partial observations .
Revised: We refer to such partially informative views as partial observations .

Original: The difficulty with partial observations occurs because distances computed between representations must include contributions from all parts of the representation vector.
Revised: The difficulty with partial observations occurs because distances computed between representations must include contributions from all parts of the representation vector.

Original: If the views provided are diverse, and therefore contain partially disjoint features, their representations may appear different to a naive distance metric.
Revised: If the views provided are diverse, and therefore contain partially disjoint features, their representations may appear different to a naive distance metric.

Original: For example, two puzzle pieces may contain different information about the whole picture.
Revised: For example, two puzzle pieces may contain different information about the whole picture.

Original: We call this the problem of integrative representation learning , where we wish to obtain a representation that integrates different but overlapping information from each element of a set.
Revised: We call this the problem of integrative representation learning , where we wish to obtain a representation that integrates different but overlapping information from each element of a set.

Original: In this paper, we provide a probabilistic formalism for a few-shot objective that is able to learn to capture representations in partially observable settings.
Revised: In this paper, we provide a probabilistic formalism for a few-shot objective that is able to learn to capture representations in partially observable settings.

Original: It does so by building on a product of experts (Hinton, 2002) to utilise representation uncertainty: a high variance in a representation component indicates that the given view of the data poorly informs the given component, while low variance indicates it informs it well.
Revised: It does so by building on a product of experts (Hinton, 2002) to utilise representation uncertainty: a high variance in a representation component indicates that the given view of the data poorly informs the given component, while low variance indicates it informs it well.

Original: Given multiple views of the data, the product of experts component in
Revised: Given multiple views of the data, the product of experts component in

Original: POEM combines the representations, weighting by the variance, to get a maximally informative and consistent representation from the views.
Revised: POEM combines the representations, weighting by the variance, to get a maximally informative and consistent representation from the views.

Original: To comprehensively evaluate our approach, we adapt a large-scale few-shot learning benchmark,
Revised: To comprehensively evaluate our approach, we adapt a large-scale few-shot learning benchmark,

Original: Meta-Dataset (Triantafillou et al., 2020), to evaluate representation learning from partial observations.
Revised: Meta-Dataset (Triantafillou et al., 2020), to evaluate representation learning from partial observations.

Original: We demonstrate that our approach, Partial Observation Experts Modelling (POEM), is able to outperform standard few-shot baselines on our adapted benchmark,
Revised: We demonstrate that our approach, Partial Observation Experts Modelling (POEM), is able to outperform standard few-shot baselines on our adapted benchmark,

Original: Partially Observed MetaDataset (PO-Meta-Dataset), while still matching state-of-the-art on the standard benchmark.
Revised: Partially Observed MetaDataset (PO-Meta-Dataset), while still matching state-of-the-art on the standard benchmark.

Original: Finally, we demonstrate the potential for our approach to be applied to meta-learn representations of environments from the partial views observed by an agent exploring that environment.
Revised: Finally, we demonstrate the potential for our approach to be applied to meta-learn representations of environments from the partial views observed by an agent exploring that environment.

Original: The main contributions of this work are: 1) A probabilistic formalism, POEM, that enables representation learning under partial observability; 2) Comprehensive experimental evaluation of POEM on an adaptation of Meta-Dataset designed to evaluate representation learning under partial observability, demonstrating that this approach outperforms standard baselines in this setting while still matching state-of-the-art on the standard fully observed benchmark; 3) A demonstration of a potential application of POEM to meta-learn representations of environments from partial observations.
Revised: The main contributions of this work are: 1) A probabilistic formalism, POEM, that enables representation learning under partial observability; 2) Comprehensive experimental evaluation of POEM on an adaptation of Meta-Dataset designed to evaluate representation learning under partial observability, demonstrating that this approach outperforms standard baselines in this setting while still matching state-of-the-art on the standard fully observed benchmark; 3) A demonstration of a potential application of POEM to meta-learn representations of environments from partial observations.

Original:  R ELATED W ORK 2. C ONTRASTIVE L EARNING
Revised:  R ELATED W ORK 2. C ONTRASTIVE L EARNING

Original: Contrastive learning extracts features that are present in multiple views of a data item, by encouraging representations of related views to be close in an embedding space (Ericsson et al., 2021).
Revised: Contrastive learning extracts features that are present in multiple views of a data item, by encouraging representations of related views to be close in an embedding space (Ericsson et al., 2021).

Original: In computer vision and natural language applications these views typically consist of different augmentations of data items, which are carefully crafted to preserve semantic features, and thereby act as an inductive bias to encourage the contrastive learner to retain these consistent features (Le-Khac et al., 2020).
Revised: In computer vision and natural language applications these views typically consist of different augmentations of data items, which are carefully crafted to preserve semantic features, and thereby act as an inductive bias to encourage the contrastive learner to retain these consistent features (Le-Khac et al., 2020).

Original: A challenge in this approach is to prevent representational ‘collapse’, where all views are mapped to the same representation.
Revised: A challenge in this approach is to prevent representational ‘collapse’, where all views are mapped to the same representation.

Original: Standard contrastive approaches such as Contrastive Predictive
Revised: Standard contrastive approaches such as Contrastive Predictive

Original: Coding (Oord et al., 2019), MoCo (He et al., 2020), and SimCLR (Chen et al., 2020) handle this by computing feature space distance measures relative to the distances for negative views – pairs of views that are encouraged to be distinct in the embedding space.
Revised: Coding (Oord et al., 2019), MoCo (He et al., 2020), and SimCLR (Chen et al., 2020) handle this by computing feature space distance measures relative to the distances for negative views – pairs of views that are encouraged to be distinct in the embedding space.

Original: Other approaches avoid the need for negatives to prevent collapse by using stop-gradients (BYOL, Grill et al.
Revised: In this work we take a similar approach, where the negative views are partial observations of distinct items, but we aim to learn to unify features from differing views, not just retain the consistent features.

Original: SimSiam, Chen & He (2020)), clustering (SwAV, Caron et al. , whitening (Ermolov et al., 2020) or regularisation (Barlow Twins Zbontar et al. (2021); VICReg Bardes et al. (2022)).
Revised: We learn to learn a contrastive representation from partial views.

Original: State-of-the-art representation learning approaches such as CLIP (Radford et al., 2021) leverage contrastive learning across modalities to match corresponding image and text representations, using a softmax matching probability of the same form as those used in embedding-based meta-learners (Vinyals et al., 2016).
Revised: We note that state-of-the-art representation learning approaches such as CLIP (Radford et al., 2021), which leverage contrastive learning across modalities, also suffer from extracting only a limited subset of features (F¨urst et al., 2022) due to using an embedding-based approach (Vinyals et al., 2016) to match image and text representations.

Original: 2. E MBEDDING -B ASED M ETA -L EARNING
Revised: 2. E MBEDDING -B ASED M ETA -L EARNING

Original: Embedding-based meta-learners similarly learn representations of classes by extracting features that are consistently present in the data samples (generally referred to as shots in the meta-learning literature) provided for each class, such that the class of new samples can be identified with a similarity measure (Hospedales et al., 2020).
Revised: Embedding-based meta-learners similarly learn representations of classes by extracting features that are consistently present in the data samples (generally referred to as shots in the meta-learning literature) provided for each class, such that the class of new samples can be identified with a similarity measure (Hospedales et al., 2020).

Original: In effect, the other classes form the negatives in a discriminative contrastive task.

Original: These methods generally differ in terms of their approach to combine features, and the distance metric used.
Revised: These methods generally differ in terms of their approach to combine features, and the distance metric used.

Original: Prototypical Networks (Snell et al., 2017) use a Euclidian distance between the query representation and the average over the support representations for a class (referred to as a prototype ).
Revised: Prototypical Networks (Snell et al., 2017) use a Euclidian distance between the query representation and the average over the support representations for a class (referred to as a prototype ).

Original: Relation Networks (Sung et al., 2018) use the same prototype representation as Prototypical Networks, but use a parameterised relation module to learn to compute the similarity between the query sample and the prototype rather than using a Euclidian distance.
Revised: Relation Networks (Sung et al., 2018) use the same prototype representation as Prototypical Networks, but use a parameterised relation module to learn to compute the similarity between the query sample and the prototype rather than using a Euclidian distance.

Original: Matching

Original:  Networks (Vinyals et al., 2016) use a Cosine distance between the query sample and each support sample as a weighting over the support labels, and so perform few-shot classification without unifying the support representations.
Revised: Matching Networks (Vinyals et al., 2016) use a Cosine distance between the query sample and each support sample as a weighting over the support labels, and so perform few-shot classification without unifying the support representations.

Revised: None of these approaches are designed to unify partially informative support samples.

Original: The approach closest to that proposed in this paper is by (Edwards & Storkey, 2017), where the authors map the different views to a statistic with an associated covariance through a variational approach.
Revised: The approach closest to that proposed in this paper is by (Edwards & Storkey, 2017), where the authors map the different views to a statistic with an associated covariance through a variational approach.

Original: However there is no control of the contribution of each view to the variance, and the covariance is diagonal, so the approach is  unsuitable for partial observation.
Revised: However there is no control of the contribution of each view to the variance, and the covariance is spherical, so the approach is also unsuitable for partial observation.

Original: 2. O PTIMISATION -B ASED M ETA -L EARNING
Revised: 2. O PTIMISATION -B ASED M ETA -L EARNING

Original: The few-shot classification task can also be solved withouta contrastive representation learning approach.
Revised: The few-shot classification task can also be solved without learning embeddings.

Original: One sensible baseline, fine-tuning of a previously pre-trained large model, simply treats each few-shot task as a standard classification problem (Nakamura & Harada, 2019).
Revised: One sensible baseline, fine-tuning of a previously pre-trained large model, simply treats each few-shot task as a standard classification problem (Nakamura & Harada, 2019).

Original: For each task, one or more additional output layers are added on top of a pre-trained embedding network and trained to predict the classes of the support set (alongside optionally finetuning the embedding network).
Revised: For each task, one or more additional output layers are added on top of a pre-trained embedding network and trained to predict the classes of the support set (alongside optionally finetuning the embedding network).

Original: This can then be utilised to predict the classes of the query set.
Revised: This can then be utilised to predict the classes of the query set.

Original: Taking this approach a step further, Model-Agnostic Meta-Learning (MAML) (Finn et al., 2017) learns the initialisation of the embedding network, such that it can be rapidly fine-tuned on a new few-shot task.
Revised: Taking this approach a step further, Model-Agnostic Meta-Learning (MAML) (Finn et al., 2017) learns the initialisation of the embedding network, such that it can be rapidly fine-tuned on a new few-shot task.

Original: At each step, the additional output layers are trained to predict the classes of the support set.

Original: The gradient from the query set loss is then backpropagated into the embedding network (requiring second-order gradients) to learn the optimal initialisation of the embedding function (and optionally the output layers) to achieve maximal accuracy on the query sets after training on the corresponding support sets.

Original: Given the generality of this approach, many variants of this method now exist, such as MAML++, Antoniou et al.
Revised: Given the generality of this approach, many variants of this method now exist, such as MAML++, Antoniou et al.

Original: Meta-SGD (Li et al., 2017), CAVIA (Zintgraf et al., 2019) and fo-Proto-MAML (Triantafillou et al., 2020).
Revised: Meta-SGD (Li et al., 2017), CAVIA (Zintgraf et al., 2019) and fo-Proto-MAML (Triantafillou et al., 2020).

Revised: One variant, LEO (Rusu et al., 2019), performs the meta-optimisation on a latent representation of the embedding parameters, learned using a relational network (Sung et al., 2018).

Revised: However, none of these variants of this fundamental optimisation based approach to few-shot learning (referred to as ’MAML’ for the remainder of this work) have a mechanism for integrating partial information from the entire support set for comparison with a partial query observation at inference time.

Original: 2. F ACTOR A NALYSIS AND P RODUCT OF E XPERTS
Revised: 2. O THER M ETA -L EARNING A PPROACHES

Original: Factor analysis is the linear counterpart to modern representation learners, but where partial observability is inherently expressed in  the model.
Revised: Probabilisitic meta-learning methods, such as VERSA (Gordon et al., 2019), DKT (Patacchiola et al., 2020) and Amortised Bayesian Prototype Meta-Learning (Sun et al., 2021), often unify both embedding-based and optimisation based meta-learning by learning to output a posterior distribution that captures uncertainty in predictions, but do not use uncertainty in features to optimally combine support set information. Other recent work, such as DeepEMD (Zhang et al., 2022), has considered the use of attention mechanisms or transformers with image patches (Hiller et al., 2022; Dong et al., 2020), or augmentations (Chen et al., 2021a).

Original: Each view can be thought of as expressing a different linear combination of latent Gaussian factors; at inference time, each view informs each  of the latent factors a different amount.
Revised: However, the purpose of these approaches is to identify informative patches or features within each support example, to improve fine-grained few-shot learning performance or interpretability where relevant features may occupy only a small region of the samples. As far as the authors are aware, there are no existing meta-learning methods that aim to integrate partial information from across the support set for comparison with a partially informative query. 2. P ARTIAL O BSERVABILITY AND P RODUCT OF E XPERTS Factor analysis is the linear counterpart to modern representation learners, but where partial observability is inherently expressed in the model.

Original: The inferential model for the latent space in factor analysis is a product of each of the conditional Gaussian factors.
Revised: The inferential model for the latent space in factor analysis is a product of each of the conditional Gaussian factors.

Original: In general, this form of inferential model can be captured as a product of experts (Hinton, 2002).
Revised: In general, this form of inferential model can be captured as a product of experts (Hinton, 2002).

Original: When those experts are Gaussian distributions (Williams et al., 2001), this product of experts is fully tractable.
Revised: When those experts are Gaussian distributions (Williams et al., 2001), this product of experts is fully tractable.

Original: By focusing on the inferential components rather than the linear model, it is possible to generalise factor analysis inference to nonlinear mappings.
Revised: By focusing on the inferential components rather than the linear model, it is possible to generalise factor analysis inference to nonlinear mappings (Tang et al., 2012).

Original: Hence the product of experts can be used more flexibly and easily than directly building deep factor analysis with neural network models (Tang et al., 2012) when only an inferential component is required (as with representation learning).
Revised: However, when only an inferential component is required (as with representation learning), the product of experts can be used more flexibly, as in our approach below.

Original:  T HEORETICAL F ORMALISM
Revised:  T HEORETICAL F ORMALISM

Original: In this section, we introduce POEM, which incorporates a product of experts model for combining different views with some prior representation, and then uses that representation to evaluate a query viewpoint.
Revised: In this section, we introduce POEM, which incorporates a product of experts model for combining different views with a prior representation, and then uses that representation to classify a query view.

Original: 3. P RODUCT OF E XPERT P ROTOTYPES
Revised: 3. P RODUCT OF E XPERT P ROTOTYPES

Original: Let each data item consist of a number of views,with each view denoted x v , indexed by v = 1 , 2 , . . , V , where the number ofviews V might vary between one data item and another. Let a superscript index count over data items, so x nv is the v th view of an n th data item. Furthermore, we
Revised: Let us consider data corresponding to partial observations, or views, of a set of items.

Original: In common with most few-shot frameworks, we arrange  data into support sets and query sets.
Revised: In common with most few-shot frameworks, we arrange the data into support sets and query sets.

Original: Each support set consists of m data items:
Revised: Each support set consists of M data items:

Original: S =
Revised: S =

Original: X m | m = 1 , 2 , . .
Revised: X m | m = 1 , 2 , . .

Original: , M } , where the m th item X 
Revised: , M } , where the m th item X m collects V m views, where V may vary with m .

Original:  m collects
Revised: Let x mv denote the v th view of the m th data item, such that

Original: V m = dim( X m ) views.
Revised: X m = { x m 1 , x m 2 , . . . , x mV m } .

Original: The items in the support set are sampled randomly from the training dataset.
Revised: The items in the support set are sampled randomly from the training dataset.

Original: The query point consists of a single different view corresponding to one and only one of the M items in the support set. We denote the query point by x ∗ .
Revised: The query point, denoted x ∗ , here consists of a single different view corresponding to one and only one of the M items in the support set (although in general we may consider N query points simultaneously).

Original: We cast our representation learning problem as a meta-learning task.
Revised: We cast our representation learning problem as a meta-learning task.

Original: We must learn a  representation derived from the support set that can be compared with a representation from a query set.
Revised: We must learn a unified representation derived from the support set that can be compared with a representation of the query view.

Original: We want that comparison to enable us to infer which support set item m = m ∗ the query view belongs to.
Revised: We want that comparison to enable us to infer which support set item m = m ∗ the query view belongs to.

Original: In this paper we are concerned with partial observability ; that is, not every data view will inform the whole representation.
Revised: In this paper we are concerned with partial observability ; that is, not every data view will inform the whole representation.

Original: So instead of mapping each view to a deterministic point, we map each view to a factor – a normalised density that indicates the uncertainty in each component.
Revised: So instead of mapping each view to a deterministic point representation, we map each view to a distributional representation where each component is a normalised density that indicates the uncertainty in that component (called a factor ).

Original: We denote this factor ϕ .
Revised: We denote this conditional density ϕ , and on implementation parameterise the parameters of the distribution ϕ with a neural network.

Original: We combine the  factors for each view together using a product of experts, which integrates a prior distribution along with the different views such that views with low variance in a component strongly inform that component.
Revised: We combine the corresponding factors for each view together using a product of experts, which integrates a prior distribution along with the different views such that views with low variance in a component strongly inform that component.

Original: For a given support set, we compute a product of experts distribution for the representation z m :
Revised: For a given support set, we compute a product of experts distribution for the representation z m :

Original: [Equation]
Revised: [Equation]

Original: where p ( z ) is a prior density over the latent space.
Revised: where p ( z ) is a prior density over the latent space.

Original: Now for a query point with a view that matches other views from e.g. data item m , we can use Bayes rule to compute the probability that the query point would be generated from the corresponding representation z m by
Revised: Now for a query point with a view that matches other views from e [Equation] g. data item m , we can use Bayes rule to compute the probability that the query point would be generated from the corresponding representation z m by

Original: [Equation]
Revised: [Equation]

Original: where, again, p ( z ) =
Revised: where, again, p ( z ) =

Original: (cid:82) d
Revised: (cid:82) d

Original: x p ( x ) ϕ
Revised: x p ( x ) ϕ

Original: ( z | x ) is the prior.
Revised: ( z | x ) is the prior.

Original: We put Eq.2 and Eq.1 together and marginalise over z m to get the marginal predictive distribution
Revised: We put Eq.2 and Eq.1 together and marginalise over z m to get the marginal predictive distribution

Original: [Equation]
Revised: [Equation]

Original: where
Revised: where

Original: [Equation]
Revised: [Equation]

Original: and
Revised: and

Original: [Equation]
Revised: [Equation]

Original: The marginal predictive p ( x ∗ | X m ) is used to form the training objective.
Revised: The marginal predictive p ( x ∗ | X m ) is used to form the training objective.

Original: In our few shot task,  the support data is chosen and then the query view is chosen uniformly at random to matchthe views of one of the support data items.
Revised: In our few shot task, we wish to maximize the likelihood for the correct match of query point to support set, accumulated across all support/query selections indexed with t from the dataset.

Original: Let the hypothesis H m indicate the event that the query view x ∗ comes from support point m . Then
Revised: This provides a complete negative log marginal likelihood objective to be minimized, as derived in appendix A.2:

Original: [Equation]
Revised: [Equation]

Original: From this we can formulate the training task: we wish to maximize the likelihood for the correct match of query point to support set, accumulated across all support/query selections from the dataset. Denote the t th support set by S t , the t th query point by x ∗ t , and let m t denote the support point with views that match the view of the query point. Then the complete negative log marginal likelihood objectiveto be minimized is [Equation]
Revised: Full pseudocode for training POEM with this objective is provided in appendix A.3.

Original: 3. I NTERPRETATION OF O BJECTIVE
Revised: 3. I NTERPRETATION OF O BJECTIVE

Original: While the normalised factors ϕ can be chosen from any distribution class, we take ϕ to be Gaussian  for the remainder of this paper, rendering the normalisation integrals in Eq.
Revised: While the normalised factors ϕ can be chosen from any distribution class, we take ϕ to be Gaussian with parameterised mean and precision for the remainder of this paper, rendering the integral in Eq.

Original: 18 analytic.
Revised: 5 analytic.

Original: Approximating the prior p ( z ) by a Gaussian also renders Eq.
Revised: Approximating the prior p ( z ) by a Gaussian also renders Eq.

Original: 19 analytic.
Revised: 6 analytic.

Original: We note that other distributions with analytic products, such as Beta distributions, may also be of interest in certain applications, but we leave an investigation of othersuch distributional forms for ϕ for further work.
Revised: We note that other distributions with analytic products, such as Beta distributions, may also be of interest in certain applications, but we leave an investigation of other distributional forms for ϕ to further work.

Original: If the representations from each view for a support point are aligned with each other and the query view (the means of all the Gaussians are similar), they will have a greater overlap and the integral of the resulting product of Gaussians will be larger, leading to a greater value of λ ( y, X ) .
Revised: If the representations from each view for a support point are aligned with each other and the query view (the means of all the Gaussians are similar), they will have a greater overlap and the integral of the resulting product of Gaussians will be larger, leading to a greater value of λ ( y, X ) .

Original: Furthermore, increasing the precisions for aligned Gaussian components leads to greater λ ( y, X ) , while, up to a limit, decreasing the precisions for non-aligned Gaussian components leads to greater λ ( y, X ) .
Revised: Furthermore, increasing the precisions for aligned Gaussian components leads to greater λ ( y, X ) , while, up to a limit, decreasing the precisions for non-aligned Gaussian components leads to greater λ ( y, X ) .

Original: While the numerator in Eq. 4, λ ( y, X ) , quantifies the overlap of the support set with the query, the denominator λ ′ ( X ) contrasts this with the overlap of the support set representation with the prior.
Revised: While the numerator in Eq. 4, λ ( y, X ) , quantifies the overlap of the support set with the query, the denominator λ ′ ( X ) contrasts this with the overlap of the support set representation with the prior.

Original: Together, this factor is enhanced if it is beneficial in overlap terms to replace the prior with the query representation, and reduced if such a replacement is detrimental.
Revised: Together, this factor is enhanced if it is beneficial in overlap terms to replace the prior with the query representation, and reduced if such a replacement is detrimental.

Original: A greater consistency between query and combined support set representations intuitively leads to a greater probability that the query belongs to the class of the corresponding support set, effectively extending Prototypical
Revised: A greater consistency between query and combined support set representations intuitively leads to a greater probability that the query belongs to the class of the corresponding support set, effectively extending Prototypical

Original: Networks to a probabilistic latent representation space (Snell et al., 2017).
Revised: Networks to a probabilistic latent representation space (Snell et al., 2017).

Original: As a result, this objective is a generalisation of a Prototypical Network that allows for (a) learnable weighted averaging over support examples based on their informativeness to a given component; (b) learnable combinations of features that only appear in subsets ofthe support examples (via differing relative precisions of components within support representations), and (c) partial comparison of the query sample with the support samples (via differing relative precisions within the query sample).
Revised: As a result, this objective is a generalisation of a Prototypical Network that allows for (a) learnable weighted averaging over support examples based on their informativeness to a given component; (b) learnable combinations of features from subsets of support examples (via differing relative precisions of components within support representations), and (c) partial comparison of the query sample with the support samples (via differing relative precisions within the query).

Original: With all precisions fixed to 1, this approach reproduces Prototypical Networks, neglecting small differences in scaling factors that arise with varying numbers of views.
Revised: With all precisions fixed to 1, this approach reproduces Prototypical Networks, neglecting small differences in scaling factors that arise with varying numbers of views.

Original: This relationship is derived in Appendix A.2.
Revised: This relationship is derived in Appendix A.4.

Original:  E XPERIMENTAL E VALUATION
Revised:  E XPERIMENTAL E VALUATION

Original: There is a lack of established benchmarks specifically targeted at the evaluation of representation learning under partial observability.
Revised: There is a lack of established benchmarks specifically targeted at the evaluation of representation learning under partial observability.

Original: To design a comprehensive benchmark for few-shot representation learning under partial observability, we leverage Meta-Dataset (Triantafillou et al., 2020), a recently proposed collection of few-shot learning benchmarks.
Revised: To design a comprehensive benchmark for few-shot representation learning under partial observability, we leverage Meta-Dataset (Triantafillou et al., 2020), a recently proposed collection of few-shot learning benchmarks.

Original: We selected Meta-Dataset as the basis for our adapted benchmark as it consists of diverse datasets involving natural, human-made and text-based visual concepts, with a variety of fine-grained classification tasks that require learning from varying and unbalanced numbers of samples and classes.
Revised: We selected Meta-Dataset as the basis for our adapted benchmark as it consists of diverse datasets involving natural, human-made and text-based visual concepts, with a variety of fine-grained classification tasks that require learning from varying and unbalanced numbers of samples and classes.

Original: As a result, our derived benchmark inherits these properties to provide a robust measure of the ability of a learning approach to learn representations from partial observations.
Revised: As a result, our derived benchmark inherits these properties to provide a robust measure of the ability of a learning approach to learn representations from partial observations.

Original: To extend Meta-Dataset to incorporate partial observability, we take multiple views of each sample and divide these views into support and query sets.
Revised: To extend Meta-Dataset to incorporate partial observability, we take multiple views of each sample and divide these views into support and query sets.

Original: Our adapted few-shot classification task is to predict which sample a query view comes from, given a selection of support views of that sample, as demonstrated in Figure 2.
Revised: Our adapted few-shot classification task is to predict which sample a query view comes from, given a selection of support views of that sample, as demonstrated in Figure 2.

Original: In keeping with the spirit of Meta-Dataset, we vary the number of ways in the task (now the number of images) from 5 to 25, taken from between 1 to 5 classes.
Revised: In keeping with the spirit of Meta-Dataset, we vary the number of ways in the task (now the number of images) from 5 to 25, taken from between 1 to 5 classes.

Original: Views are generated by applying the standard augmentation operations used in SimCLR (Chen et al., 2020) and most other selfsupervised learning methods.
Revised: Views are generated by applying the standard augmentation operations used in SimCLR (Chen et al., 2020) and most other selfsupervised learning methods.

Original: However, to emphasise the focus on partial observability, the size of  the random crops and the number of views was fixed, such that the entire support set for a sample contains a maximum of 50% of the image.
Revised: However, to emphasise the focus on partial observability, the size of [Figure] the random crops and the number of views was fixed, such that the entire support set for a sample contains a maximum of 50% of the image.

Original: We also maintain a constant number of query views per sample.
Revised: We also maintain a constant number of query views per sample.

Original: Viewpoint information consisting of the coordinates of the view is provided to make it possible for learners to understand where a view fits into a representation even in the absence of overlapping views.
Revised: Viewpoint information consisting of the coordinates of the view is provided to make it possible for learners to understand where a view fits into a representation even in the absence of overlapping views.

Original: Full details of the definition of the task are provided in appendix A.3.
Revised: Full details of the definition of the task are provided in appendix A.5.

Original: [Figure]

Original: We apply our proposed evaluation procedure to all datasets included in Meta-Dataset with a few exceptions.
Revised: We apply our proposed evaluation procedure to all datasets included in Meta-Dataset with a few exceptions.

Original: ILSVRC (ImageNet, Russakovsky et al.
Revised: ILSVRC (ImageNet, Russakovsky et al.

Original: (2015)) was not included since our network backbones were pre-trained on this dataset, including the standard few-shot test classes (which is also why this dataset was subsequently removed from the updated benchmark, MetaDataset-v2 (Dumoulin et al., 2021)).
Revised: (2015)) was not included since our network backbones were pre-trained on this dataset, including the standard few-shot test classes (which is also why this dataset was subsequently removed from the updated benchmark, MetaDataset-v2 (Dumoulin et al., 2021)).

Original: Traffic Signs (Stallkamp et al., 2011) and MSCOCO (Lin et al., 2015) were not included since these datasets are fully reserved for evaluation by Meta-Dataset and so do not have a training set specified (and dataset transfer is not in the scope of this work).
Revised: Traffic Signs (Stallkamp et al., 2011) and MSCOCO (Lin et al., 2015) were not included since these datasets are fully reserved for evaluation by Meta-Dataset and so do not have a training set specified.

Original: Quick Draw (FernandezFernandez et al., 2019) was also not included since this dataset was found to be too large to use within the memory constraints of standard RTX2080 GPUs.
Revised: Quick Draw (Fernandez-Fernandez et al., 2019) was also not included since this dataset was found to be too large to use within the memory constraints of standard RTX GPUs.

Original: This leaves six diverse datasets: Aircraft (Maji et al., 2013),
Revised: This leaves six diverse datasets: Aircraft (Maji et al., 2013),

Original: Birds (Wah et al., 2011), Flowers (Nilsback & Zisserman, 2008), Fungi (Schroeder, Brigit, 2018), Omniglot (Lake et al., 2015) and Textures (Cimpoi et al., 2014), on all of which our models were trained, validated and tested on according to the data partitions specified by the Meta-Dataset benchmark.
Revised: Birds (Wah et al., 2011), Flowers (Nilsback & Zisserman, 2008), Fungi (Schroeder, Brigit, 2018), Omniglot (Lake et al., 2015) and Textures (Cimpoi et al., 2014), on all of which our models were trained, validated and tested on according to the data partitions specified by the Meta-Dataset benchmark.

Original: The resulting benchmark, Partially Observed Meta-Dataset (PO-Meta-Dataset), therefore requires that the learner coherently combine the information from the support views into a consistent representation of the sample, such that the query view can be matched to the sample it originated from.
Revised: The resulting benchmark, Partially Observed Meta-Dataset (PO-Meta-Dataset), therefore requires that the learner coherently combine the information from the support views into a consistent representation of the sample, such that the query view can be matched to the sample it originated from.

Original: Since a maximum of 50% of each sample is seen in the support set, the task also requires generalisation to correctly match and classify query views.
Revised: Since a maximum of 50% of each sample is seen in the support set, the task also requires generalisation to correctly match and classify query views.

Original: 4. I MPLEMENTATION D ETAILS
Revised: 4. I MPLEMENTATION D ETAILS

Original: We utilise a re-implementation of Meta-Dataset benchmarking in PyTorch (Paszke et al., 2019) which closely replicates the Meta-Dataset sampling procedure of uniformly sampling classes, followed by a balanced query set (since all classes are considered equally important) and unbalanced support sets (to mirror realistic variations in the appearances of classes).
Revised: We utilise a re-implementation of Meta-Dataset benchmarking in PyTorch (Paszke et al., 2019) which closely replicates the Meta-Dataset sampling procedure of uniformly sampling classes, followed by a balanced query set (since all classes are considered equally important) and unbalanced support sets (to mirror realistic variations in the appearances of classes).

Original: The experimental implementation, including full open-source code and data will be available on publication.
Revised: The experimental implementation, including full open-source code and data will be available on publication.

Original: Following the MD-Transfer procedure used in Meta-Dataset (and in the subsequent Meta-Datasetv2 Dumoulin et al. (2021)), we leverage a single ResNet-18 (He et al., 2015)
Revised: Following the MD-Transfer procedure used in Meta-Dataset, we leverage a single ResNet-18 (He et al., 2015)

Original: classifier pre-trained on ImageNet (Russakovsky et al., 2015) at 126 × 126 resolution.
Revised: classifier pre-trained on ImageNet (Russakovsky et al., 2015) at 126 × 126 resolution.

Original: Since both a mean and precision must be learned to fully specify the model ϕ v ( z | x nv ) , we add two simple 3-layer MLP heads onto this backbone for POEM, each maintaining an embedding size of512 .
Revised: Since both a mean and precision must be learned to fully specify the model ϕ v ( z | x nv ) , we add two simple 3-layer MLP heads onto this backbone for POEM, each maintaining an embedding size of

Original: For fair comparison, we also add the same 3-layer MLP head onto the backbone for the baselines.
Revised: For fair comparison, we also add the same 3-layer MLP head onto the backbone for the baselines.

Original: Using a larger embedding for the baselines was not found to be beneficial.
Revised: Using a larger embedding for the baselines was not found to be beneficial.

Original: During training, gradients are backpropagated through the entire network such that both the randomly initialised heads and pre-trained backbones are learned/fine-tuned.
Revised: During training, gradients are backpropagated through the entire network such that both the randomly initialised heads and pre-trained backbones are learned/fine-tuned.

Original: We use a  selection ofthe standard meta-learning baselines utilised by Meta-Dataset for our reimplementation, including both representation-based (Prototypical Networks, Snell et al. and  optimisation-based (Finetuning Nakamura & Harada (2019); MAML, Finn et al.
Revised: We use a representative selection of meta-learning baselines utilised by Meta-Dataset for our reimplementation. This includes a strong naive baseline (Finetuning, Nakamura & Harada (2019)), an embedding-based approach (Prototypical Network, Snell et al. and an optimisation-based approach (MAML, Finn et al.

Original: (2017)) approaches.
Revised: (2017)), all modernised to use the ResNet-18 backbone as described above. Recent competitions, such as the NeurIPS 2021 MetaDL Challenge (Baz et al., 2022; 2021), have demonstrated that these fundamental approaches, updated to use modern pre-trained backbones that are finetuned on the meta-task (exactly as in our experiments below) are still generally stateof-the-art for novel datasets (Chen et al., 2021b), and so form strong baselines.

Original: This re-implementation enables us to ensure that all  comparisons between learners are fair, utilising the same benchmark parameters, model architectures and where applicable, hyperparameters.
Revised: In addition, our re-implementation enables us to ensure that all learners are optimised for Meta-Dataset and that comparisons between learners are fair, utilising the same benchmark parameters, model architectures and where applicable, hyperparameters.

Original: Crucially, given the close connection between POEM and Prototypical Networks, we ensure that all hyperparameters, including learning rates, scheduling and architectures are identical.
Revised: Crucially, given the close connection between POEM and Prototypical Networks, we ensure that all hyperparameters, including learning rates, scheduling and architectures are identical for both methods.

Original: 4. R ESULTS
Revised: 4. R ESULTS

Original: Our results on our proposed representation learning benchmark, PO-Meta-Dataset, are provided in table 1.
Revised: Our results on this novel representation learning benchmark, PO-Meta-Dataset, are given in table 1.

Original: [Table]
Revised: [Table]

Original: observations. The results show that POEM outperforms the baselines at identifying views of images across a diverse range of datasets, demonstrating the benefits of the approach to learn and match representations from partial observations.
Revised: The results show that POEM outperforms the baselines at identifying views of images across a diverse range of datasets, demonstrating the benefits of the approach to learn and match representations from partial observations.

Original: The only exception is the Textures dataset, for which the finetuning baseline performs particularly strongly.
Revised: The only exception is the Textures dataset, for which the finetuning baseline performs particularly strongly.

Original: We hypothesise that this is because the images in the Textures dataset are relatively uniform compared to the other datasets, so capturing the relative location of views is less important than identifying very fine grained features that distinguish the samples, which optimisation-based approaches are particularly effective at.
Revised: We hypothesise that this is because the images in the Textures dataset are relatively uniform compared to the other datasets, so capturing the relative location of views is less important than identifying very fine grained features that distinguish the samples, which optimisation-based approaches are particularly effective at.

Original: 4. A BLATION : M ETA -D ATASET
Revised: 4. A BLATION : M ETA -D ATASET

Original: To demonstrate that the observed benefit of POEM over the baselines is due to the requirement of the task to learn coherent representations from partial observations, we also evaluate our approach against the baselines on the established Meta-Dataset benchmark.
Revised: To demonstrate that the observed benefit of POEM over the baselines is due to the requirement of the task to learn coherent representations from partial observations, we also evaluate our approach against the baselines on the established Meta-Dataset benchmark.

Original: We now follow the standard fewshot learning procedure as described in the paper [Table] but keep all learners identical to those used in the evaluation above
Revised: We now follow the standard fewshot learning procedure as described in the paper (Triantafillou et al., 2020), but keep all learners identical to those used in the evaluation above.

Original: [Table] Our results on the standard Meta-Dataset benchmark are providedbelow in table 2.[Table] head.
Revised: Our results on the standard Meta-Dataset benchmark are provided in table 2.

Original: POEM is comparable with the baselines.

Original: As expected, we find that POEM performs comparably with the baselines.
Revised: As expected, we find that POEM performs comparably with the baselines.

Original: Although Meta-Dataset provides realistic few-shot learning tasks in terms of diversity of visual concepts, fine-grained classes and variable shots and ways, each sample generally contains complete information including all relevant features for the visual concept in question.
Revised: Although Meta-Dataset provides realistic few-shot learning tasks in terms of diversity of visual concepts, fine-grained classes and variable shots and ways, each sample generally contains complete information including all relevant features for the visual concept in question.

Original: Correctly classifying query samples does not generally require any unification of views from support examples, but simply the identification of common features.
Revised: Correctly classifying query samples does not generally require any unification of views from support examples, but simply the identification of common features.

Original: As a result, we see that the additional capacity of POEM to learn to weight support examples and combine partial features does not provide a significant performance improvement over the baselines at few-shot classification in this fully observable benchmark.
Revised: As a result, we see that the additional capacity of POEM to learn to weight support examples and [Table]

Original: In support of our hypothesis that feature uncertainty is not useful on this benchmark, we find that the variance in the precisions relative to the means output by the POEM model generally decreases during training and becomes negligible for all datasets, indicating that the precisions are not being utilised to improve performance and that the POEM objective is reducing to the Prototypical
Revised: In support of our hypothesis that feature uncertainty is not useful on this benchmark, we find that the variance in the precisions relative to the means output by the POEM model generally decreases during training and becomes negligible for all datasets, indicating that the precisions are not being utilised to improve performance and that the POEM objective is reducing to the Prototypical

Original: Network objective, as discussed in section 3.2.
Revised: Network objective, as discussed in section 3.2.

Original: This is further evidenced by the very similar performances of POEM and the Prototypical Network across the entire set of datasets.
Revised: This is further evidenced by the very similar performances of POEM and the Prototypical Network across the entire set of datasets.

Original: However, on
Revised: However, on

Original: PO-Meta-Dataset, we find that the relative variance in the precisions to the means is much larger on convergence, which leads to the improved performance of POEM over the Prototypical Network observed in Table 1.
Revised: PO-Meta-Dataset, we find that the relative variance in the precisions to the means is much larger on convergence, which leads to the improved performance of POEM over the Prototypical Network observed in Table 1.

Original: This is shown in appendix A.4.
Revised: This is shown in appendix A.6.

Original:  D EMONSTRATION OF L EARNING R EPRESENTATIONS OF E NVIRONMENTS
Revised:  D EMONSTRATION OF L EARNING R EPRESENTATIONS OF E NVIRONMENTS

Original: We now apply POEM to the equivalent task of learning a representation of an environment from the partial observations collected by an agent exploring that environment.
Revised: We now apply POEM to the equivalent task of learning a representation of an environment from the partial observations collected by an agent exploring that environment.

Original: To do so, we utilise the 2D gridworld environment, MiniGrid (Chevalier-Boisvert et al., 2018).
Revised: To do so, we utilise the 2D gridworld environment, MiniGrid (Chevalier-Boisvert et al., 2018).

Original: We consider the 11 × 11 Simple Crossing environment, which consists of a procedurally generated maze where the agent is required to traverse from the top left corner to the goal in the bottom right corner.
Revised: We consider the 11 × 11 Simple Crossing environment, which consists of a procedurally generated maze where the agent is required to traverse from the top left corner to the goal in the bottom right corner.

Original: The MiniGrid environment provides an agent-centric viewpoint at each step in the trajectory, consisting of a 7 × 7 window of the environment in front of the agent, taking into account the agent’s current direction and where the line of sight is blocked by walls.
Revised: The MiniGrid environment provides an agent-centric viewpoint at each step in the trajectory, consisting of a 7 × 7 window of the environment in front of the agent, taking into account the agent’s current direction and where the line of sight is blocked by walls.

Original: 5. M ETA -L EARNING E NVIRONMENT R EPRESENTATIONS VIA F EW -S HOT C LASSIFICATION
Revised: 5. M ETA -L EARNING E NVIRONMENT R EPRESENTATIONS VIA F EW -S HOT C LASSIFICATION

Original: To generate few-shot episodes, we utilise two agents: an optimal agent that takes the optimal trajectory from the start to the goal, and an exploratory agent that is incentivised to explore all possible views in the environment.
Revised: To generate few-shot episodes, we utilise two agents: an optimal agent that takes the optimal trajectory from the start to the goal, and an exploratory agent that is incentivised to explore all possible views in the environment.

Original: The support set for each environment is generated by running the optimal agent in the environment and collecting the partial observations of this agent at each step during its trajectory.
Revised: The support set for each environment is generated by running the optimal agent in the environment and collecting the partial observations of this agent at each step during its trajectory.

Original: The query set is similarly generated by running the exploratory agent in the environment, filtering out any observations that are contained within the support set, and then randomly sampling the desired number of queries from the remaining observations.
Revised: The query set is similarly generated by running the exploratory agent in the environment, filtering out any observations that are contained within the support set, and then randomly sampling the desired number of queries from the remaining observations.

Original: We generate these few-shot episodes dynamically, and train POEM tolearn to combine the support samples (partial observations from the optimal trajectory) into a representation of the environment, such that it can classify which environment a novel query observation has been collected from.
Revised: We generate these few-shot episodes dynamically, and train POEM to combine the support samples (partial observations from the optimal trajectory) into a representation of the environment, such that it can classify which environment a novel query observation has been collected from.

Original: Sample observations of a set of  environments and observations from those environments are shown in figures 3 and 4.
Revised: A set of sample environments and observations from those environments are shown in figures 3 and 4.

Original: [Figure] [Figure] All observations are provided as pixels to a standard convolutional backbone, with the corresponding agent location and direction appended to this representation and passed through an MLP head, equivalent to the procedure utilised for the adapted Meta-Dataset experiments.
Revised: All observations are provided as pixels to a standard convolutional backbone, with the corresponding agent location and direction appended to this representation and passed through an MLP head, equivalent to the procedure utilised for the adapted Meta-Dataset experiments.

Original: As a baseline comparison, we also train a Prototypical Network with an identical architecutre on this task.
Revised: As a baseline comparison, we also train a Prototypical Network with an identical architecture on this task.

Revised: Additionally, we train an equivalent recurrent network architecture typically applied to POMDP tasks such as this (Hausknecht & Stone, 2017), by adding a GRU layer (Cho et al., 2014; Chung et al., 2014) where the hidden state of the GRU is updated at each timestep and then extracted as the unified representation of the agent.

Original: We find that POEM trains more quickly and reaches almost 10% higher final environment recognition performance than  the Prototypical Network  over 100 test episodes (81.1% vs 72.4%), as shown in appendix A.6.
Revised: We find that POEM trains more quickly and reaches almost 10% higher final environment recognition performance than both the Prototypical Network and GRU-based approach over 100 test episodes (81.1% vs 72.4% and 72.1%), as shown in appendix A.8.

Original: This is a result of POEM’s capacity to associate each observation with only part of the representation, rather than simply average observation representations.
Revised: This is a result of POEM’s capacity to associate each observation with only part of the representation.

Original: 5. R ECONSTRUCTING E NVIRONMENTS FROM P ARTIAL O BSERVATION T RAJECTORIES
Revised: 5. R ECONSTRUCTING E NVIRONMENTS FROM P ARTIAL O BSERVATION T RAJECTORIES

Original: Having learned an environment encoder using the few-shot learning procedure above, we now investigate the extent to which our representations can be used to regenerate the environment.
Revised: Having learned an environment encoder using the few-shot learning procedure above, we now investigate the extent to which our representations can be used to reconstruct the environment.

Original: As above, wedynamically generate trajectories with the optimal agent and feed these through theenvironment encoder to generate a representation of the environment.
Revised: As above, we generate trajectories with the optimal agent and feed these through the encoder to generate a representation of the environment.

Original: An MLP decoder is then trained to reconstruct the original environment layout from the learned environment representation.
Revised: An MLP decoder is then trained to reconstruct the original environment layout from the learned environment representation.

Original: The decoder attempts to predict a one-hot representation of each possible grid cell, with a mean squared error loss.
Revised: The decoder attempts to predict a one-hot representation of each possible grid cell, with a mean squared error loss.

Original: Given the trained encoder and decoder, this now allows us to generate a map of the environment the optimal agent has traversed solely from the agent’s partial observations,  without ever having seen the environment as a whole.
Revised: Given the trained encoder and decoder, we are now able to generate a map of the environment the optimal agent has traversed, solely from the agent’s partial observations, and without ever having seen the environment as a whole.

Original: A sample of environments alongside their reconstructions are shown in figure 5.
Revised: A sample of environments alongside their reconstructions are shown in figure 5.

Original:  We see that the reconstructions clearly capture the approximate structure of each environment, demonstrating that the agent has been able to integrate its observations from along its trajectory into a single consistent representation.
Revised: vations of the agent. We see that the reconstructions clearly capture the approximate structure of each environment, demonstrating that the agent has been able to integrate its observations from along its trajectory into a single consistent representation.

Original: Since POEM enables the representation to be updated incrementally with each partial observation of the environment at inference time, it would be possible for an agent to update an internal environment representation at each step in its trajectory.
Revised: Since POEM enables the representation to be updated incrementally with each partial observation of the environment at inference time, it would be possible for an agent to update an internal environment representation at each step in its trajectory.

Original: There is potential for utilising this approach for learninga representation of an environment  to be beneficial in the context of exploration for reinforcement learning, but we leave such an investigation to future work.
Revised: There is potential for utilising this approach for learning environment representations to be beneficial in the context of exploration for reinforcement learning, but we leave such an investigation to future work.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: In this work, we have introduced Partial Observation Experts Modelling (POEM), a contrastive meta-learning approach for few-shot learning in partially-observable settings.
Revised: In this work, we have introduced Partial Observation Experts Modelling (POEM), a contrastive meta-learning approach for few-shot learning in partially-observable settings.

Original: Unlike other standard contrastive and embedding-based meta-learning approaches POEM utilises representational uncertainty to enable observations to inform only part of a representation vector
Revised: Unlike other standard contrastive and embedding-based meta-learning approaches, POEM utilises representational uncertainty to enable observations to inform only part of a representation vector.

Original: This probabilistic formalism enables consistent representation learning from multiple observations with a few-shot learning objective
Revised: This probabilistic formalism enables consistent representation learning from multiple observations with a few-shot learning objective.

Original: We have demonstrated that

Original:  POEM is comparable to the state of -the -art base - lines on a comprehensive few-shot learning benchmark, and outperforms these baselines when this benchmark is adapted to evaluate representation learning from partial observations.
Revised: We have demonstrated that POEM is comparable to the state-of-the-art baselines on a comprehensive few-shot learning benchmark, and outperforms these baselines when this benchmark is adapted to evaluate representation learning from partial observations.

Original: We have also demonstrated a promising potential application for POEM to learn representations of an environ - ment from an agent’s partial observations.
Revised: We have also demonstrated a promising potential application for POEM to learn representations of an environment from an agent’s partial observations.

Original: We hope that this research inspires further work into the challenging task of learning representations under partial observability and the creation of more realistic partial observability benchmarks
Revised: We hope that this research inspires further work into the challenging task of learning representations under partial observability and the creation of more realistic partial observability benchmarks. R EPRODUCIBILITY S TATEMENT

Original: A A PPENDIX A. G AUSSIAN P RODUCT R ULES
Revised: All results reported in this paper are reproducible. The experimental implementations, including full open-source code will be made available upon acceptance for publication. E THICS S TATEMENT This work considers the problem of learning representations from partial observations.

Revised: The authors are not aware of any specific ethical concerns associated with improving the quality of learned representations from partial observations, other than those already associated with any work in the field of machine learning

Original: Assuming the latent variable model ϕv ( zn | xnv ; θ ) to be a diagonal covariance multivariate Gaussian, the resulting integrals over latent variables become integrals over Gaussian products.
Revised: Assuming the latent variable model ϕ ( z | x ) to be a diagonal covariance multivariate Gaussian, the resulting integrals over latent variables become integrals over Gaussian products.

Original: This allows both λ ( y , X ) (Equation 18) and λ ′ ( X ) (Equation 19) in the marginal predictive distribution (Equation 4) to be evaluated analytically using the following univariate Gaussian product rules on each independent dimension (Roweis, Sam, 1999).
Revised: This allows both λ ( y , X ) (Equation 5) and λ ′ ( X ) (Equation 6) in the marginal predictive distribution (Equation 4) to be evaluated analytically using the following univariate Gaussian product rules on each independent dimension (Roweis, Sam, 1999).

Original: Since a product of Gaussians (cid:81) i N ( µ i , τ − 1 i ) is itself a Gaussian, we have (cid:81) i N ( µ i , τ − 1 i ) =
Revised: Since a product of Gaussians (cid:81) i N ( µ i , τ − 1 i ) is itself a Gaussian, we have (cid:81) i N ( µ i , τ − 1 i ) =

Original: SN ( µ, τ − 1 ) , where
Revised: SN ( µ, τ − 1 ) , where

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: Therefore the integral of a Gaussian product is given by  S .
Revised: Therefore the integral of a Gaussian product is given by the resulting normalisation S .

Revised: In the case of evaluating the marginal predictive distribution p ( x ∗ | X m ) (equation 4), this gives S ∗ S

Original: The quotient of two Gaussian pdfs N (µ + , τ + − 1) /N
Revised: S ′ S = S ∗ S ′ where S is the normalisation constant of the support product, S ∗ is the normalisation constant of the product of the query and normalised support product, and S ′ is the normalisation constant of the product of the prior p ( z ) and normalised support product. In reality, p ( z ) generally has little impact as it is typically flat ( τ → 0 ) over the region of non-negligible density of the product (cid:81) Vv =1 ϕ

Original: ( µ − , τ − 1 − ) with τ − < τ + is also Gaussian, given{+and so S ′ ≈ 
 and we find S ∗ S ′ ≈ S ∗ so the ratios λλ ′ in the objective can be approximated+} by S Q N ( µ Q , τ Q ) where
Revised: ( z | x v ) {+and so S ′ ≈ 
 and we find S ∗ S ′ ≈ S ∗ so the ratios λλ ′ in the objective can be approximated+} by S ∗ , as in the simplified pseduocode in appendix A.3.

Revised: A. D ERIVATION OF O BJECTIVE FROM M ARGINAL P REDICTIVE D ISTRIBUTION

Revised: In section 3, we derived the marginal predictive distribution:

Original: [Equation]
Revised: [Equation]

Revised: (11) where

Original: [Equation]
Revised: [Equation]

Revised: and

Original: [Equation] 
Revised: [Equation] [Equation]

Revised: In our few shot task, the support data is chosen and then the query view is chosen uniformly at random to match the views of one of the support data items.

Revised: Let the hypothesis H m indicate the event that the query view x ∗ comes from support point m .

Revised: Then

Original: [Equation]
Revised: [Equation]

Original: For the case of computing the marginal predictive distribution (Equation 4), we have a Gaussian product in both the numerator (with product precision, mean and normalisation τ + , µ + , S + ) and in the denominator (with product precision, mean and normalisation τ − , µ − , S − ), allowing us to compute the quotient normalisation S Q , and yielding the resulting total integral as:
Revised: From this we can formulate the training task: we wish to maximize the likelihood for the correct match of query point to support set, accumulated across all support/query selections from the dataset. Denote the t th support set by S t , the t th query point by x ∗ t , and let m t denote the support point with views that match the view of the query point. Then the complete negative log marginal likelihood objective to be minimized is:

Original: [Equation]
Revised: [Equation]

Revised: A. P SEUDOCODE

Revised: [Table]

Revised: Observation Experts Modelling

Revised: Require:

Revised: Training meta-set D train ∈ T

Revised: Require:

Revised: Learning rate α 1: Initialise dual-headed network ϕ θ ( z | x ) 2: ▷ Heads correspond to mean µ and precision τ of Gaussian embedding z 3: while not converged do 4:

Revised: Sample task instance T i =

Revised: ( X , x ∗ ) ∼ D train 5: ▷ Support set X consists of V m views of item m ∈ { 1 , ..., M } .

Revised: 6: ▷ Query set x ∗ consists of N queries, each one view from any one item.

Revised: Encode each view in support set X into Gaussian z using ϕ

Revised: ( z | X ) 8:

Revised: Encode each query view in x ∗ into Gaussian z ∗ using ϕ

Revised: ( z ∗ | x ∗ ) 9: for m ∈ { 1 , ..., M } do 10:

Revised: Compute Gaussian product over views (cid:81) V m v =1 ϕ

Revised: ( z | x mv ) (using results in A.1) 11:

Revised: ▷ This gives unified support representation (global environment representation) 12: for n ∈ { 1 , ..., N } do 13:

Revised: Compute Gaussian product of query with support product ϕ

Revised: ( z ∗ n | x ∗ n ) (cid:81) V m v =1 ϕ

Revised: ( z | x mv ) 14: end for 15: end for 16:

Revised: Normalise resulting query-support normalisation constants S m n =

Revised: S mn (cid:80) m S mn across items 17:

Revised: Compute negative log of S m ∗ n for correct support as loss L ( { D t } , { x ∗ t } ) (eq. 7) 18: ▷ Negative log likelihood for correct support 19:

Revised: Perform gradient step w.r.t.

Revised: θ : θ ← ϕ − α ∇ θ L ( { D t } , { x ∗ t } ) 20: end while

Original: A. E QUIVALENCE OF P ROTOTYPICAL N ETWORK O BJECTIVE TO POEM O BJECTIVE WITH F IXED P RECISIONS
Revised: A. E QUIVALENCE OF P ROTOTYPICAL N ETWORK O BJECTIVE TO POEM O BJECTIVE WITH F IXED P RECISIONS

Original: The probability of a query q belonging to class k using the POEM objective is given by:
Revised: The probability of a query x ∗ belonging to class n using the POEM objective is given by:

Original: [Equation]
Revised: [Equation]

Original: as defined in equation 8, where
Revised: as defined in equation [Equation] where

Original: [Equation]
Revised: [Equation]

Original: Taking the precisions of the all Gaussian factors ϕ in λ and λ ′ to be 1, we can apply the Gaussian product rules given in appendix A.1 to calculate λ and λ ′ analytically.
Revised: Taking the precisions of the all Gaussian factors ϕ in λ and λ ′ to be [Equation] we can apply the Gaussian product rules given in appendix A.1 to calculate λ and λ ′ analytically.

Original: We find that this gives:
Revised: We find that this gives:

Revised: [Equation]

Revised: [Equation]

Revised: where µ is the representation mean of the query

Original: [Equation]
Revised: [Equation]

Original: where µ is the representation mean of the query, and µ ni is the representation mean of support sample i for class n , and V n is the number of support samples for class n .
Revised: and µ ni is the representation mean of support sample i for class n [Equation] and V n is the number of support samples for class n .

Original: Equivalently, the probability of a query with representation vector µ belonging to a class n using a
Revised: Equivalently, the probability of a query with representation vector µ belonging to a class n using a

Original: Prototypical Network objective is given by:
Revised: Prototypical Network objective is given by:

Original: [Equation] [Equation]
Revised: [Equation] [Equation]

Original: We find that these are equivalent aside from the scaling factors
Revised: We find that these are equivalent aside from the scaling factors

Original: V m (2)( V m +1) which only have a (significant) effect when there are varying numbers of samples by class, and a greater effect when the number of samples is smaller.
Revised: V m [Equation] [Equation] V m +1) which only have a (significant) effect when there are varying numbers of samples by class, and a greater effect when the number of samples is smaller.

Original: Experimentally, we find that these scaling factors make little difference, as demonstrated in table 2 of section 4.3.
Revised: Experimentally, we find that these scaling factors make little difference, as demonstrated in table 2 of section 4.3.

Original: A. PO-M ETA -D ATASET B ENCHMARK A DDITIONAL D ETAILS
Revised: A. PO-M ETA -D ATASET B ENCHMARK A DDITIONAL D ETAILS

Original: Parameters used for adapted PO-Meta-Dataset.
Revised: Parameters used for adapted PO-Meta-Dataset are provided in Table A.5.

Original: All parameters not listed chosen to match MetaDataset defaults.
Revised: All parameters not listed chosen to match Meta-Dataset defaults.

Original: All augmentations are applied using Torchvision, with parameters specified.
Revised: All augmentations are applied using Torchvision, with parameters specified.

Original: [Table]

Original: [Table]

Original: [Table]

Original: 2) , p ( apply ) =

Original: [Table]
Revised: [Table]

Original: ((3 , 3) , (1 , 0 , 2 . 0)) , p ( apply ) = 0 .

Original: All results computed over three runs.
Revised: All results computed over three runs.

Original: The Finetuning, Prototypical Network and POEM baselines were run on on-premise RTX2080 GPUs.
Revised: The Finetuning, Prototypical Network and POEM baselines were run on on-premise RTX2080 GPUs.

Original: MAML required more memory and compute than available, so was run on cloud A100s.
Revised: MAML required more memory and compute than available, so was run on cloud A100s.

Original: A. R ELATIVE V ARIANCE OF P RECISIONS D URING T RAINING ON M ETA -D ATASET AND M ETA -M ETA -D ATASET
Revised: A. R ELATIVE V ARIANCE OF P RECISIONS D URING T RAINING ON M ETA -D ATASET AND M ETA -M ETA -D ATASET

Original: The plot below shows the evolution of the variance in the representation precisions relative to the variance in the representation means learned by POEM on two distinct datasets, Aircraft and VGG
Revised: The plot below shows the evolution of the variance in the representation precisions relative to the variance in the representation means learned by POEM on two distinct datasets, Aircraft and VGG

Original: Flowers.
Revised: Flowers.

Original: We see that for standard few-shot learning on Meta-Dataset, the variance in precisions is negligible relative to the variance in the means, demonstrating that the representational uncertainty is not useful in this task.
Revised: We see that for standard few-shot learning on Meta-Dataset, the variance in precisions is negligible relative to the variance in the means, demonstrating that the representational uncertainty is not useful in this task.

Original: Meanwhile, we see the variance in the precisions relative to the variance in the means becoming large before converging to a value of O (100) on the Meta-Meta-Dataset task, demonstrating that learning relative precisions is useful in this setting since each support sample only informs part of the representation.
Revised: Meanwhile, we see the variance in the precisions relative to the variance in the means becoming large before converging to a value of O (100) on the Meta-Meta-Dataset task, demonstrating that learning relative precisions is useful in this setting since each support sample only informs part of the representation.

Original: A. L EARNING R EPRESENTATIONS OF E NVIRONMENTS A DDITIONAL D ETAILS
Revised: A. L EARNING R EPRESENTATIONS OF E NVIRONMENTS A DDITIONAL D ETAILS

Original: [Table]
Revised: Additional details about the parameters used for learning environment representations from agent observations are provided in Table

Original: [Table]
Revised: [Table]

Original:  E NVIRONMENT R ECOGNITION A CCURACY D URING T RAINING
Revised: A. E NVIRONMENT R ECOGNITION A CCURACY D URING T RAINING

Revised: POEM trains more quickly on the environment recognition task and reaches a higher final performance than an equivalent Prototypical Network or Recurrent Network (GRU) (81.1% vs 72.4% and 72.1% ) over a subsequent 100 test episodes

Original: [Table]
Revised: [Table]

Original: trains more quickly on the environment recognition task and reaches a higher final performance (81.1% vs 72.4% ) over a subsequent 100 test episodes
Revised: [Figure] [Table] [Figure]

Original: [Table]
Revised: [Table]

Revised: Environment Recognition Accuracy

Revised: [Table]