Group Equivariant Conditional Neural Processes

1 INTRODUCTION . Data symmetry has played a significant role in the deep neural networks . In particular , a convolutional neural network , which play an important part in the recent achievements of deep neural networks , has translation equivariance that preserves the symmetry of the translation group . From the same point of view , many studies have aimed to incorporate various group symmetries into neural networks , especially convolutional operation ( Cohen et al. , 2019 ; Defferrard et al. , 2019 ; Finzi et al. , 2020 ) . As example applications , to solve the dynamics modeling problems , some works have introduced Hamiltonian dynamics ( Greydanus et al. , 2019 ; Toth et al. , 2019 ; Zhong et al. , 2019 ) . Similarly , Quessard et al . ( 2020 ) estimated the action of the group by assuming the symmetry in the latent space inferred by the neural network . Incorporating the data structure ( symmetries ) into the models as inductive bias , can reduce the model complexity and improve model generalization . In terms of inductive bias , meta-learning , or learning to learn , provides a way to select an inductive bias from data . Meta-learning use past experiences to adapt quickly to a new task T ∼ p ( T ) sampled from some task distribution p ( T ) . Especially in supervised meta-learning , a task is described as predicting a set of unlabeled data ( target points ) given a set of labeled data ( context points ) . Various works have proposed the use of supervised meta-learning from different perspectives ( Andrychowicz et al. , 2016 ; Ravi & Larochelle , 2016 ; Finn et al. , 2017 ; Snell et al. , 2017 ; Santoro et al. , 2016 ; Rusu et al. , 2018 ) . In this study , we are interested in neural processes ( NPs ) ( Garnelo et al. , 2018a ; b ) , which are meta-learning models that have encoder-decoder architecture ( Xu et al. , 2019 ) . The encoder is a permutation-invariant function on the context points that maps the contexts into a latent representation . The decoder is a function that produces the conditional predictive distribution of targets given the latent representation . The objective of NPs is to learn the encoder and the decoder , so that the predictive model generalizes well to new tasks by observing some points of the tasks . To achieve the objective , an NP is required to learn the shared information between the training tasks T , T ′ ∼ p ( T ) : the data knowledge Lemke et al . ( 2015 ) . Each task T is represented by one dataset , and multiple datasets are provided for training NPs to tackle a meta-task . For example , we consider a meta-task that completing the pixels that are missing in a given image . Often , images are taken by the same condition in each dataset , respectively . While the datasets contain identical subjects of images ( e.g. , cars or apples ) , the size and angle of the subjects in the image may be different ; the datasets have group symmetry , such as scaling and rotation . Therefore , it is expected that pre-constraining NPs to have group equivariance improves the performance of the NPs at those datasets . In this paper , we investigate the group equivalence of NPs . Specifically , we try to answer the following two questions , ( 1 ) can NPs represent equivariant functions ? ( 2 ) can we explicitly induce the group equivariance into NPs ? In order to answer the questions , we introduce a new family of NPs , EquivCNP , and show that EquivCNP is a permutation-invariant and group-equivariant function theoretically and empirically . Most relevant to EquivCNP , ConvCNP ( Gordon et al. , 2019 ) shows that using general convolution operation leads to the translation equivariance theoretically and experimentally ; however it does not consider incorporation of other groups . First , we introduce the decomposition theorem for permutation-invariant and group-equivariant maps . The theorem suggests that the encoder maps the context points into a latent variable , which is a functional representation , in order to preserve the data symmetry . Thereafter , we construct EquivCNP by following the theorem . In this study , we adopt LieConv ( Finzi et al. , 2020 ) to construct EquivCNP for practical implementation . We tackle a 1D synthetic regression task ( Garnelo et al. , 2018a ; b ; Kim et al. , 2019 ; Gordon et al. , 2019 ) to show that EquivCNP with translation equivariance is comparable to conventional NPs . Furthermore , we design a 2D image completion task to investigate the potential of EquivCNP with several group equivariances . As a result , we demonstrate that EquivCNP enables zero-shot generalization by incorporating not translation , but scaling equivariance . 2 RELATED WORK . 2.1 NEURAL NETWORKS WITH GROUP EQUIVARIANCE . Our works build upon the recent advances in group equivariant convolutional operation incorporated into deep neural networks . The first approach is group convolution introduced in ( Cohen & Welling , 2016 ) , where standard convolutional kernels are used and their transformation or the output transformation is performed with respect to the group . This group convolution induces exact equivariance , but only to the action of discrete groups . In contrast , for exact equivariance to continuous groups , some works employ harmonic analysis so as to find the basis of equivariant functions , and then parameterize convolutional kernels in the basis ( Weiler & Cesa , 2019 ) . Although this approach can be applied to any type of general data ( Anderson et al. , 2019 ; Weiler & Cesa , 2019 ) , it is limited to local application to compact , unimodular groups . To address these issues , LieConv ( Finzi et al. , 2020 ) and other works ( Huang et al. , 2017 ; Bekkers , 2019 ) use Lie groups . Our EquivCNP chooses LieConv to manage group equivariance for simplicity of the implementation . There are several works that study deep neural networks using data symmetry . In some works , in order to solve machine learning problems such as sequence prediction or reinforcement learning , neural networks attempt to learn a data symmetry of physical systems from noisy observations directly ( Greydanus et al. , 2019 ; Toth et al. , 2019 ; Zhong et al. , 2019 ; Sanchez-Gonzalez et al. , 2019 ) . While both these studies and EquivCNP can handle data symmetries , EquivCNP is not limited to specific domains such as physics . Furthermore , Quessard et al . ( 2020 ) let the latent space into which neural networks map data , have group equivariance , and estimated the parameters of data symmetries . In terms of using group equivariance in the latent space , EquivCNP is similar to this study but differs from being able to use various group equivariance . 2.2 FAMILY OF NEURAL PROCESSES . NPs ( Garnelo et al. , 2018a ; b ) are deep generative models for regression functions that map an input xi ∈ Rdx into an output yi ∈ Rdy . In particular , given an arbitrary number of observed data points ( xC , yC ) : = { ( xi , yi ) } Ci=1 , NPs model the conditional distribution of the target value yT at some new , unobserved target data point xT , where ( xT , yT ) : = { ( xj , yj ) } Tj=1 . Fundamentally , there are two NP variants : deterministic and probabilistic . Deterministic NPs ( Garnelo et al. , 2018a ) , known as conditional NPs ( CNPs ) , model the conditional distribution as : p ( yT |xT , xC , yC ) : = p ( yT |xT , rC ) , where r represents a function that maps data sets ( xC , yC ) into a finite-dimensional vector space in a permutation-invariant way and rC : = r ( xC , yC ) ∈ Rd is the feature vector . The function r can be implemented by DeepSets ( Zaheer et al. , 2017 ) . The likelihood p ( yT |xT , rC ) is modeled by Gaussian distribution factorized across the targets ( xj , yj ) with mean and variance of prediction { ( xj , yj ) } Tj=1 by passing inputs rC and xj through the MLP . The CNP is trained by maximizing the likelihood . Probabilistic NPs include a latent variable z . The NP infers q ( z|rC ) given an input rC using the reparametrization trick ( Kingma & Welling , 2013 ) and models such a conditional distribution as : p ( yT |xT , xC , yC ) : = ∫ p ( yT |xT , rC , z ) q ( z|rC ) dz and it is trained by maximizing an ELBO : L ( φ , θ ) = Ez∼qφ ( z|xT , yT ) [ log pθ ( yT |xT ) ] − KL [ qφ ( z|xT , yT ) ‖pθ ( z|xC , yC ) ] . NPs have various useful properties : i ) Scalability : the computational cost of NPs scales as O ( n+m ) with respect to n contexts and m targets of data , ii ) Flexibility : NPs can define a conditional distribution of an arbitrary number of target points , conditioning an arbitrary number of observations , iii ) Permutation invariance : the encoder of NPs uses Deepsets ( Zaheer et al. , 2017 ) to make the target prediction permutation invariant . Thanks to these properties , Galashov et al . ( 2019 ) replace Gaussian processes in Bayesian optimization , contextual multi-armed bandit , and Sim2Real tasks . While there are many NP variants ( Kim et al. , 2019 ; Louizos et al. , 2019 ; Xu et al. , 2019 ) to improve the performance of NPs , those do not take group equivariance into account yet . The most similar to EquivCNP , ConvCNP ( Gordon et al. , 2019 ) incorporated only translation equivariance . In contrast , EquivCNP can incorporate not only translation but also other groups such as rotation and scaling . 3 DECOMPOSITION THEOREM . In this section , we consider group convolution . We first prepare some definition and teminology . Let X and Y ⊂ R be the input space and output space , respectively . We define ZM = ( X × Y ) M as a collection of M input-output pairs , Z≤M = ⋃M n=1Zn as the collection of at most M pairs , and Z = ⋃∞ m=1Zm as the collection of finitely many pairs . Let [ n ] = { 1 , . . . , n } for n ∈ N , and let Sn be the permutation group on [ n ] . The action of Sn on Zn is defined as πZn : = ( ( xπ−1 ( 1 ) , yπ−1 ( 1 ) ) , . . . , ( xπ−1 ( n ) , yπ−1 ( n ) ) ) , where π ∈ Sn and Zn ∈ Zn . We define the multiplicity of Zn = ( ( x1 , y1 ) , . . . , ( xn , yn ) ) ∈ Zn by mult ( Zn ) : = sup { | { i ∈ [ n ] : xi = x̂ } | : x̂ = x1 , . . . , xn } and the multiplicity of Z ′ ⊆ Z by mult ( Z ′ ) : = supZn∈Z′ mult ( Zn ) . Then , a collection Z ′ ⊆ Z is said to have multiplicity K if mult ( Z ′ ) = K. Mathematically , symmetry is described in terms of group action . The following group equivariant maps represent to preserve the symmetry in data . Definition 1 ( Group Equivariance and Invariance ) . Suppose that a group G acts on sets S and S ′ . Then , a map Φ : S → S ′ is called G-equivariant when Φ ( g · s ) = g ·Φ ( s ) holds for arbitrary g ∈ G and s ∈ S. In particular , when G acts on S ′ trivially ( i.e. , g · s′ = s′ for g ∈ G and s′ ∈ S ′ ) , the G-equivariant map is said to be G-invariant : Φ ( g · s ) = Φ ( s ) . Then , we can derive the following theorem , which decompose a permutation-invariant and group equivariant function into two tractable functions . Note that this theorem has been proved by Gordon et al . ( 2019 ) when G is a translation group . Theorem 2 ( Decomposition Theorem ) . Let G be a group . Let Z ′≤M ⊆ ( X ×Y ) ≤M be topologically closed , permutation-invariant and G-invariant with multiplicity K. For a function Φ : Z ′≤M → Cb ( X , Y ) , the following conditions are equivalent : Algorithm 1 Prediction of Group Equivariant Conditional Neural Process Input : ρ =LieConv , RBF kernel ψ , context { xi , yi } Ni=1 , target { x∗j } Mj=1 lower , upper← range ( ( xi ) Ni=1 ∪ ( xj ) Mj=1 ) ( tk ) T k=1 ← uniform_grid ( lower , upper ; γ ) // Encoding the context information into representation h ( i.e . Encoder ) h← ∑N i=1 φK+1 ( yi ) ψ ( [ x ∗ j , tk ] , xi ) h. ( µj , Σj ) > = LieConvNet ( h ) ( x∗j ) // Decoder Output : { ( µj , Σj ) } Mj=1 ( I ) Φ is continuous , permutation-invariant and G-equivariant . ( II ) There exist a function spaceH and a continuous G-equivariant function ρ : H → Cb ( X , Y ) and a continuous G-invariant interpolating kernel ψ : X 2 → R such that Φ ( Z ) = ρ ( m∑ i=1 φK+1 ( yi ) ψxi ) where φK+1 : Y → RK+1 is defined by φK+1 ( y ) : = [ 1 , y , y2 , . . . , yK ] > . Thanks to the Theorem 2 , we can construct the permutation-invariant and group-equivariant NPs whose form of encoder and decoder is determined . In this paper , we call Φ as EquivDeepSet .

Summary: This paper presents EquivCNP which extends Conditional Neural Processes (CNP) to incorporate symmetries of the data, e.g. rotation and scaling. The approach utilizes a combination of LieConv (Finzi et al., 2020)  and DeepSet (Zaheer et al., 2017) to achieve the equivariance in the data space and permutation invariance across the samples in a dataset. They provide empirical results on a 1D regression task with synthetic and 2D image completion tasks using digital clock digits dataset which they constructed.