Learning to Generate 3D Shapes with Generative Cellular Automata

1 INTRODUCTION . Probabilistic 3D shape generation aims to learn and sample from the distribution of diverse 3D shapes and has applications including 3D contents generation or robot interaction . Specifically , learning the distribution of shapes or scenes can automate the process of generating diverse and realistic virtual environments or new object designs . Likewise , modeling the conditional distribution of the whole scene given partial raw 3D scans can help the decision process of a robot , by informing various possible outputs of occluded space . The distribution of plausible shapes in 3D space is diverse and complex , and we seek a scalable formulation of the shape generation process . Pioneering works on 3D shape generation try to regress the entire shape ( Dai et al . ( 2017 ) ) which often fail to recover fine details . We propose a more modular approach that progressively generates shape by a sequence of local updates . Our work takes inspiration from prior works on autoregressive models in the image domains , such as the variants of pixelCNN ( van den Oord & Kalchbrenner ( 2016 ) ; van den Oord et al . ( 2016 ; 2017 ) ) , which have been successful in image generation . The key idea of pixelCNN ( van den Oord et al . ( 2016 ) ) is to order the pixels , and then learn the conditional distribution of the next pixel given all of the previous pixels . Thus generating an image becomes the task of sampling pixel-by-pixel in the predefined order . Recently , PointGrow ( Sun et al . ( 2020 ) ) proposes a similar approach in the field of 3D generation , replacing the RGB values of pixels with the coordinates of points and sampling point-by-point in a sequential manner . While the work proposes a promising interpretable generation process by sequentially growing a shape , the required number of sampling procedures expands linearly with the number of points , making the model hard to scale to high-resolution data . We believe that a more scalable solution in 3D is to employ the local update rules of cellular automata ( CA ) . CA , a mathematical model operating on a grid , defines a state to be a collection of cells that carries values in the grid ( Wolfram ( 1982 ) ) . The CA repeatedly mutates its states based on the predefined homogeneous update rules only determined by the spatial neighborhood of the current cell . In contrast to the conventional CA where the rules are predefined , we employ a neural network to infer the stochastic sequential transition rule of individual cells based on Markov chain . The obtained homogeneous local rule for the individual cells constitutes the 3D generative model , named Generative Cellular Automata ( GCA ) . When the rule is distributed into the group of occupied cells of an arbitrary starting shape , the sequence of local transitions eventually evolves into an instance among the diverse shapes from the multi-modal distribution . The local update rules of CA greatly reduce the search space of voxel occupancy , exploiting the sparsity and connectivity of 3D shapes . We suggest a simple , progressive training procedure to learn the distribution of local transitions of which repeated application generates the shape of the data distribution . We represent the shape in terms of surface points and store it within a 3D grid , and the transition rule is trained only on the occupied cells by employing a sparse CNN ( Graham et al . ( 2018 ) ) . The sparse representation can capture the high-resolution context information , and yet learn the effective rule enjoying the expressive power of deep CNN as demonstrated in various computer vision tasks ( Krizhevsky et al . ( 2012 ) ; He et al . ( 2017 ) ) . Inspired by Bordes et al . ( 2017 ) , our model learns sequences that are slightly different from the sampling chain but converge to the full shapes in the training data . The network successfully learns the update rules of CA , such that a single inference samples from the distribution of diverse modes along the surface . The contributions of the paper are highlighted as follows : ( 1 ) We propose Generative Cellular Automata ( GCA ) , a Markov chain based 3D generative model that iteratively mends the shape to a learned distribution , generating diverse and high-fidelity shapes . ( 2 ) Our work is the first to learn the local update rules of cellular automata for 3D shape generation in voxel representation . This enables the use of an expressive sparse CNN and reduces the search space of voxel occupancy by fully exploiting sparsity and connectivity of 3D shapes . ( 3 ) Extensive experiments show that our method has competitive performance against the state-of-the-art models in probabilistic shape completion and shape generation . 2 3D SHAPE GENERATION WITH GENERATIVE CELLULAR AUTOMATA . Let Zn be an n-dimensional uniform grid space , where n = 3 for a 3D voxel space . A 3D shape is represented as a state s ⊂ Z3 , which is an ordered set of occupied cells c ∈ Z3 in a binary occupancy grid based on the location of the surface . Note that our voxel representation is different from the conventional occupancy grid , where 1 represents that the cell is inside the surface . Instead , we only store the cells lying on the surface . This representation can better exploit the sparsity of 3D shape than the full occupancy grid . The shape generation process is presented as a sequence of state variables s0 : T that is drawn from the following Markov Chain : s0 ∼ p0 st+1 ∼ pθ ( st+1|st ) ( 1 ) where p0 is the initial distribution and pθ is the homogeneous transition kernel parameterized by θ . We denote the sampled sequence s0 → s1 → · · · → sT as a sampling chain . Given the data set D containing 3D shapes x ∈ D , our objective is to learn the parameters θ of transition kernel pθ , such that the marginal distribution of final generated sample p ( sT ) = ∑ s0 : T−1 p 0 ( s0 ) ∏ 0≤t < T pθ ( s t+1|st ) is close to the data distribution . The initial state s0 can be defined differently depending on the task to solve . For the task of probabilistic shape completion , s0 is given as the partial input shape . For shape generation , we set the initial state s0 to be the most simple state we can think of , a single cell { c } . Figure 1 presents examples of sampling chains of shape generation , where the starting shape s0 is merely a single cell . The GCA further splits the transition kernel pθ ( st+1|st ) to be the combination of local update rules on individual occupied cells ci ∈ st , as depicted in Figure 2 . The cellular transition is implemented with the sparse convolution , which is translation invariant if implemented with a fully convolutional network , and outputs the distribution of local occupied cells . Then individual predictions are aggregated by cell-wise averaging , resulting in the binary probability distribution for occupancy of each cell that follows the Bernoulli distribution : pθ ( s t+1|st ) = ∏ c∈Zn pθ ( c|st ) . ( 2 ) The next state st+1 is sampled independently for individual cells from the obtained distribution and the sampling chain continues to the next time step . For each transition pθ ( st+1|st ) , we limit the search space of the occupied cells by confining our predictions within the neighborhood of the occupied cells . The underlying assumption is that the occupied cells of a valid shape are connected and a successful generation is possible by progressive growing into the immediate neighborhood of the given state . Specifically , the output of the sparse convolution is the occupancy probability of neighborhood cells pi = pθ ( N ( ci ) |st ) , where the neighborhood cells are those that fall within a radius r ball centered at the cell , N ( ci ) = { c′ ∈ Zn| d ( ci , c′ ) ≤ r } given a distance metric d. Other cells are ignored assuming they have probability 0 of occupancy . If the input state has M occupied cells st = { c1 , · · · , cM } , the sparse convolution predicts the occupancy probability of individual cells with N -dimension vectors P = { p1 , · · · , pM } , where N is the number of neighborhood cells fixed by the distance threshold r within the uniform grid Zn . After the cell-wise averaging step , the aggregated probability is nonzero for coordinates in N ( st ) = ⋃ c∈st N ( c ) . Then the cell-wise sampling in Eq . ( 2 ) is performed only within N ( st ) , instead of the full grid Zn , leading to an efficient sampling procedure . The stochastic local transition rule pθ ( N ( ci ) |st ) changes the state of a cell ’ s immediate neighborhood N ( ci ) , but the inference is determined from a larger perception neighborhood . In contrast , classical cellular automata updates a state of a cell determined by a fixed rule given the observation of the cell ’ s immediate neighborhood . The large perception neighborhood of GCA is effectively handled by deep sparse convolutional network , and results in convergence to a single consistent global shape out of diverse possible output shapes as further discussed in the appendix F . 3 TRAINING GENERATIVE CELLULAR AUTOMATA . We aim to learn the parameters for the local transition probability pθ ( N ( ci ) |st ) , whose repetitive application generates shape that follows the complex distribution of the data set . If we have sequences of sampling chains , then all current and next states can serve as the training data . Because we only have the set of complete shapes D , we emulate the sequence of sampling chain and obtain the intermediate states . The emulated sequence of a sampling chain may start from an arbitrary state , and needs to converge to the full shape x ∈ D after local transitions to the neighbor of the previous state . A naive method would be to sample the next state st from the sampling chain and maximize pθ ( x ∩N ( st ) |st ) , the probability of the shape that is closest to x among reachable forms from the current state1 , similar to scheduled sampling ( Bengio et al . ( 2015 ) ) . While this approach clearly emulates the inference procedure , it leads to learning a biased estimator as pointed out in Huszár ( 2015 ) . More importantly , the emulated sequence can not consistently learn the full shape x as it is not guaranteed to visit the state s such that x ⊂ N ( s ) . We instead employ the use of infusion training procedure suggested by Bordes et al . ( 2017 ) . Specifically , the infusion chain , denoted as s̃0 → s̃1 → ... → s̃T , is obtained as following : s̃0 ∼ q0 ( s̃0|x ) qt ( s̃t+1|s̃t , x ) = ∏ c̃∈N ( s̃t ) ( 1− αt ) pθ ( c̃|s̃t ) + αtδx ( c̃ ) ( 3 ) where q0 indicates the initial distribution of infusion chain , and qt is the infusion transition kernel at time step t. For probabilistic shape completion s̃0 is sampled as a subset of x , while for shape generation s̃0 is a single cell c ∈ x . The transition kernel qt is defined for cells in the neighborhood c̃ ∈ N ( s̃t ) as the mixture of pθ ( c̃|s̃t ) and δx ( c̃ ) , which are the transition kernel of the sampling chain and the infusion of the ground shape x , respectively . δx ( c̃ ) is crucial to guarantee the sequence to ultimately reach the ground truth shape , and is formulated as the Bernoulli distribution with probability 1 , if c̃ ∈ x , else 0 . We set the infusion rate to increase linearly with respect to time step , i.e. , αt = max ( wt , 1 ) , where w > 0 is the speed of infusion rate as in Bordes et al . ( 2017 ) . We can prove that the training procedure converges to the shape x in the training data if it is composed of weakly connected cells . We first define the connectivity of two states . Definition 1 . We call a state s̃′ to be partially connected to state x , if for any cell b ∈ x , there is a cell c0 ∈ s̃′ ∩ x , and a finite sequence of coordinates c0 : Tb in x that starts with c0 and ends with cTb = b , where each subsequent element is closer than the given threshold distance r , i.e. , for any b ∈ x , ∃c0 : Tb , such that ci ∈ x , d ( ci , ci+1 ) ≤ r , 0 ≤ i ≤ Tb and c0 ∈ s̃′ , cTb = b . The shape x is partially connected to any s̃′ 6= ∅ if we can create a sequence of coordinates between any pair of cells in x that is composed of local transitions bounded by the distance r. This is a very weak connectivity condition , and any set that overlaps with x is partially connected to x for shapes with continuous surfaces , which include shapes in the conventional dataset . Now assuming that the state s̃t ′ is partially connected to the state x , we recursively create a sequence by defining s̃t+1 = N ( s̃t ) ∩ x , which is the next sequence of infusion chain with infusion rate 1 . Since we use a linear scheduler for infusion rate , we can assume that there exists a state s̃t ′ such that infusion rate αt ′ = 1 . The following proposition proves that the sequence ( s̃t ) t≥t′ converges to x with local transitions . Proposition 1 . Let state s̃t ′ be partially connected to state x , where x has a finite number of occupied cells . We denote a sequence of states s̃t ′ : ∞ , recursively defined as s̃t+1 = N ( s̃t ) ∩ x . Then , there exists an integer T ′ such that s̃t = x , t ≥ T ′ . Proof . The proof is found in the appendix A . The proposition states that the samples of infusion chain eventually converge to the shape x , and thus we can compute the nonzero likelihood p ( x|s̃T ) during training if T is large enough . One can also interpret the training procedure as learning the sequence of states that converges to x and is close to 1Since we defined a state as a set of occupied cells , union ( ∪ ) , intersection ( ∩ ) , and subset ( ⊂ ) operations can be defined as regular sets . Figure 3 : Qualitative comparison of probabilistic shape completion . Figure 4 : Samples from shape generation . the sampling chain . We empirically observe that most of the training samples converge to the shape x before the infusion rate becomes 1 . The training procedure can be summarized as the following : 1 . Sample infusion chain s̃0 : T from s̃0 ∼ q0 ( s̃0|x ) , s̃t+1 ∼ qt ( s̃t+1|s̃t , x ) . 2 . For each state s̃t , maximize the log-likelihood that has the closest form to x via gradient descent , i.e. , θ ← θ + η ∂ log pθ ( x∩N ( s̃ t ) |s̃t ) ∂θ . The full training algorithm utilizes a data buffer to de-correlate the gradients of subsequent time steps , and accelerates the process by controlling the time steps based on the current state . More details regarding the full training algorithm can be found in the appendix B .

Summary: The paper proposes a generative method for 3D objects (voxels representation). Given an initial voxels configuration (e.g. partial shape, or even a single voxel), the method learns a local transition kernel for a Markov chain to decide how to evolve the configuration; sampling iteratively from these probabilities leads to a final model. The paper shows results on shape completion and generation, obtaining fairly good results.