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: My understanding of the paper is that authors propose to introduce contrastive methods (Van den Oord 2019, Chen 2020) to improve the learning of stochastic processes with NN architectures. Particularly, they identify the main points where the Neural Process (NP+CNP) family of models fails, e.g. due to high-dimensionality or complex likelihood densities. The advantage of using these contrastive representations for stochastic processes is that it makes the task likelihood-free, performing better when the aforementioned problems appear.