On Contrastive Representations of Stochastic Processes

1 Introduction The stochastic process ( Doob , 1953 ; Parzen , 1999 ) is a powerful mathematical abstraction used in biology ( Bressloff , 2014 ) , chemistry ( van Kampen , 1992 ) , physics ( Jacobs , 2010 ) , finance ( Steele , 2012 ) and other fields . The simplest incarnation of a stochastic process is a random function R → R , such as a Gaussian Process ( MacKay , 2003 ) , that can be used to describe a real-valued signal indexed by time or space . Extending to random functions from R to another space , stochastic processes can model timedependent phenomena like queuing ( Grimmett and Stirzaker , 2020 ) and diffusion ( Itô et al. , 2012 ) . In meta-learning , the stochastic process can be used to describe few-shot learning tasks—mappings from images to class labels ( Vinyals et al. , 2016 ) —and image completion tasks—mappings from pixel locations to RGB values ( Garnelo et al. , 2018a ) . In computer vision , 2D views of 3D objects can be seen as observations of a stochastic process indexed by the space of possible viewpoints ( Eslami et al. , 2018 ; Mildenhall et al. , 2020 ) . Videos can be seen as samples from a time-indexed stochastic process with 2D image observations ( Zelnik-Manor and Irani , 2001 ) . ∗Equal contribution . Author ordering determined by coin flip . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Machine learning algorithms that operate on data generated from stochastic processes are therefore in high demand . We assume that we have access to only a small set of covariate–observation pairs { ( xi , yi ) Ci=1 } from different realizations of the underlying stochastic process . This might correspond to a few views of a 3D object , or a few snapshots of a dynamical system evolving in time . Whilst conventional deep learning thrives when there is a large quantity of i.i.d . data available ( Lake et al. , 2017 ) , allowing us to learn a fresh model for each realization of the stochastic process , when the context size is small it makes sense to use data from other realizations to build up prior knowledge about the domain which can aid learning on new realizations ( Reed et al. , 2018 ; Garnelo et al. , 2018a ) . Traditional methods for learning from stochastic processes , including the Gaussian Process family ( MacKay , 2003 ; Rasmussen , 2003 ) and the Neural Process family ( Garnelo et al. , 2018a , b ; Eslami et al. , 2018 ) , learn to reconstruct a realization of the process from a given context . That is , given a context set { ( xi , yi ) Ci=1 } , these methods provide a predictive distribution q ( y ? |x ? , ( xi , yi ) Ci=1 ) for the observation that would be obtained from this realization of the process at any target covariate x ? . These methods use an explicit likelihood for q , typically a Gaussian distribution . Whilst this can work well when y ? is low-dimensional and unimodal , it is a restrictive assumption . For example , when p ( y ? |x ? , ( xi , yi ) Ci=1 ) samples a high-dimensional image with colour distortion , traditional methods must learn to perform conditional image generation , a notably challenging task ( van den Oord et al. , 2016 ; Chrysos and Panagakis , 2021 ) . In this paper , we do away with the explicit likelihood requirement for learning from stochastic processes . Our first insight is that , for a range of important downstream tasks , exact reconstruction is not necessary to obtain good performance . Indeed , whilst y may be high-dimensional , the downstream target label or feature ` may be simpler . We consider two distinct settings for ` ∈ L. The first is a downstream task that depends on the covariate x ∈ X , formally a second process X → L that covaries with the first . For example , ` ( x ) could represent a class label or annotation for each video frame . The second is a downstream task that depends on the entire process realization , such as a single label for a 3D object . In both cases , we assume that we have limited labelled data , so we are in a semi-supervised setting ( Zhu , 2005 ) . To solve problems of this nature , we propose a general framework for Contrastive Representations of Stochastic Processes ( CRESP ) . At its core , CRESP consists of a flexible encoder network architecture for contexts { ( xi , yi ) C1 } that unites transformer encoders of sets ( Vaswani et al. , 2017 ; Parmar et al. , 2018 ) with convolutional encoders ( LeCun et al. , 1989 ) for observations that are images . To account for the two kinds of downstream task that may of interest , we propose a targeted variant of CRESP that learns a representations depending on the context and a target covariate x ? , and an untargeted variant that learns one representation of the context . To train our encoder , we take our inspiration from recent advances in contrastive learning ( Bachman et al. , 2019 ; Chen et al. , 2020 ) which have so far focused on representations of single observations , typically images . We define a variant of the InfoNCE objective ( van den Oord et al. , 2018 ) for contexts sampled from stochastic processes , allowing us to avoid training objectives that necessitate exact reconstruction . Rather than attempting pixel-perfect reconstruction , then , CRESP solves a self-supervised task in representation space . The CRESP framework unifies and extends recent work , building on function contrastive learning ( FCLR ) ( Gondal et al. , 2021 ) by considering targeted as well as untargeted representations and using self-attention in place of mean-pool aggregation . We develop on noise contrastive meta-learning ( Ton et al. , 2021 ) by focusing on downstream tasks rather than multi-modal reconstruction , replacing conditional mean embeddings with neural representations and using a simpler training objective . We evaluate CRESP on sinusoidal functions , 3D objects , and dynamical processes with highdimensional observations . We empirically show that our methods can handle high-dimensional observations with naturalistic distortion , unlike explicit likelihood methods , and our representations lead to improved data efficiency compared to supervised learning . CRESP performs well on a range of downstream tasks , both targeted and untargeted , outperforming existing methods across the board . Our code is publicly available at github.com/ae-foster/cresp . 2 Background . Stochastic Processes Stochastic Processes ( SPs ) are probabilistic objects defined as a family of random variables indexed by a covariate space X . For each x ∈ X , there is a corresponding random variable y|x ∈ Y living in the observation space . For example , x might represent a pose and y a photograph of an underlying object take from pose x ( see Tab . 1 ) . We assume that there is a realization F sampled from a prior p ( F ) , and that the random variable y|x is a sample from p ( y|F , x ) . Thus , for each x ∈ X , F defines a conditional distribution p ( y|F , x ) . We assume that observations are independent conditional on the realization F . Hence , the joint distribution of multiple observations at locations x1 : C from one realization of the stochastic process with prior p ( F ) is p ( y1 : C |x1 : C ) = ∫ p ( F ) C∏ i=1 p ( yi|F , xi ) dF . ( 1 ) Conversely , assuming exchangeability and consistency , the Kolmogorov Extension Theorem guarantees that the joint distribution takes the form ( 1 ) ( Øksendal , 2003 ; Garnelo et al. , 2018b ) . Neural Processes The neural process ( NP ) and conditional neural process ( CNP ) are closely related models that learn representations of data generated by a stochastic process ( SP ) 1 . The training objective for the NP and CNP is inspired by the posterior predictive distribution for SPs : given a context { ( xi , yi ) Ci=1 } , the observation at the target covariate x ? has the distribution p ( y ? |x ? , ( xi , yi ) Ci=1 ) = ∫ p ( F | ( xi , yi ) Ci=1 ) p ( y ? |F , x ? ) dF . ( 2 ) The CNP learns a neural approximation q ( y ? |x ? , ( xi , yi ) Ci=1 ) = p ( y ? |c , x ? ) to equation ( 2 ) , where c = ∑ i genc ( xi , yi ) is a permutation-invariant context representation and p ( ·|c , x ) is an explicit likelihood . Conventionally , p is a Gaussian with mean and variance given by a neural network applied to c , x . The CNP model is then trained by maximum likelihood . In the NP model , an additional latent variable u is used to represent uncertainty in the process realization , more closely mimicking ( 2 ) . A significant limitation , common to the NP family , is the reliance on an explicit likelihood . Indeed , requiring log q ( y ? |x ? , ( xi , yi ) Ci=1 ) to be large requires the model to successfully reconstruct y ? based on the context , similarly to the reconstruction term in variational autoencoders ( Kingma and Welling , 2014 ) . Furthermore , the NP objective can not be increased by extracting additional features from the context unless the predictive part of the model , the part mapping from ( c , x ) to a mean and variance , is powerful enough to use them . Contrastive Learning and Likelihood-free Inference Contrastive learning has enjoyed recent success in learning representations of high-dimensional data ( van den Oord et al. , 2018 ; Bachman et al. , 2019 ; He et al. , 2020 ; Chen et al. , 2020 ) , and is deeply connected to likelihood-free inference ( Gutmann and Hyvärinen , 2010 ; van den Oord et al. , 2018 ; Durkan et al. , 2020 ) . In its simplest form , suppose we have a distribution p ( y , y′ ) , for example y and y′ could be differently augmented versions of the same image . Rather than fitting a model to predict y′ given y , which would necessitate high-dimensional reconstruction , contrastive learning methods can be seen as learning the likelihoodratio r ( y′|y ) = p ( y′|y ) /p ( y′ ) . To achieve this , contrastive methods encode y , y′ to deterministic embeddings z , z′ , and consider additional ‘ negative ’ samples z′1 , ... , z ′ K−1 which are the embeddings of other independent samples of p ( y′ ) ( for example , taken from the same training batch as y , y′ ) . The InfoNCE training loss ( van den Oord et al. , 2018 ) is then given by LInfoNCEK = −E [ log s ( z , z′ ) s ( z , z′ ) + ∑ k s ( z , z ′ k ) ] − logK . ( 3 ) for similarity score s > 0 . Informally , InfoNCE is minimized when z is more similar to z′—the ‘ positive ’ sample—than it is to the negative samples z′1 , ... , z ′ K−1 that are independent of z . Formally , Eq . ( 3 ) is the multi-class cross-entropy loss arising from classifying the positive sample correctly . It can be shown that the optimal similarity score s is proportional to the true likelihood ratio r ( van den Oord et al. , 2018 ; Durkan et al. , 2020 ) . A key feature of InfoNCE is that learns about the predictive density p ( y′|y ) indirectly , rather than by attempting direct reconstruction .

Summary: The authors propose a method for representation of an observational context by optimization of a contrastive learning objective. In contrast to (conditional) neural process modeling, this avoids specification of a likelihood, which may be difficult and ultimately unnecessary in tasks that only require prediction of a downstream label. Experimental results illustrate the advantages of this method over related neural process models on such tasks.