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: The paper provides an extension of convolution conditional neural processes CNPs to more general Lie group equivariant CNPs. The development of the theory seems sufficiently clear to someone more familiar with the field. However, for newer readers, it seems important to be familiar with background concepts and prior work. This is not a penalizing point but rather just an observation.