Unsupervised Learning of Compositional Energy Concepts

1 Introduction . Human intelligence is characterized by its ability to learn new concepts , such as the manipulation of a new tool from only a few demonstrations [ 1 ] . Essential to this capability is the composition and re-utilization of previously learned concepts to accomplish the task at hand [ 36 ] . This is especially apparent in natural language , which is often described as a tool for making ‘ infinite use of finite means ’ [ 8 ] . Previously acquired words can be infinitely nested using a set of grammatical rules to communicate an arbitrary thought , opinion , or state one is in . In this work , we are interested in constructing a system that can discover , in an unsupervised manner , a broad set of these compositional components , as well as subsequently combine them across distinct modalities and datasets . For obtaining such decompositions , two separate lines of work exist . The first focuses on obtaining global , holistic , compositional factors by situating data points , such as human faces , in an underlying ( fixed ) factored vector space [ 24 , 48 ] . Individual factors , such as emotion or hair color , are represented as independent dimensions of the vector space , with recombination between factors corresponding to the recombination of the underlying dimensions . Due to the fixed dimensionality of the vector space , multiple instances of a single factor , such as lighting , may not be combined , nor can individual * Code and data available at https : //energy-based-model.github.io/comet/ 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . factored vector spaces from separate datasets be combined , i.e . the facial expression in an image from one dataset , and the background lighting in an image from another . To address this weakness , a separate line of work decomposes a scene into a set of underlying ‘ object ’ factors . Each object factor represents a separate set of pixels in an image , as defined by a disjoint segmentation mask [ 22 , 47 ] . Such a representation allows for the composition of individual factors by compositing the segmentation masks . However , by explicitly constraining the decomposition to be in terms of disjoint segmentation masks , relationships between individual factors of variation become more difficult to represent and capture global factors describing a scene . In this work , we propose instead to decompose a scene into a set of factors represented as energy functions . An individual energy function represents a factor by assigning low energy to scenes with said factor and high energy to scenes where said factor is absent . A scene is then generated by optimizing the sum of energies for all factors . Multiple factors can be composed together , by summing the energies for each individual factor . Simultaneously , these individual factors are defined across the entire scene , allowing energy functions to represent global factors ( lighting , camera viewpoint ) and local factors ( object existence ) . Our work is inspired by recent work [ 12 ] which shows that energy based models may be utilized to represent flexible compositions of both global and local factors . However , while [ 12 ] required supervision , as labels were used to represent concepts , we aim to decompose and discover concepts in an unsupervised manner . In our approach , we discover energy functions from separate data points by enforcing compositions of these energy functions to recompose data . Our work is further inspired by the ability of humans to effortlessly and reliably combine concepts gathered from disparate experiences . As a separate benefit of our approach , we show that we may take components inferred by one instance of our model trained on one dataset , and compose it with other components inferred by other instances of our model trained on separate datasets . We provide analysis showing why our approach , COMET† , is favorable compared to existing unsupervised approaches for decomposing scenes . We contribute the following : First , we show COMET provides a unified framework enabling us to decompose images into both global factors of variation as well as local factors of variation . Second , we show that COMET enables us to scale to more realistic datasets than previous work . Finally , we show that components obtained by COMET generalize well , and are amenable to compositions across different modes of data , and with components discovered by other instances of COMET . 2 Related Work . Our work is related to research in the areas of global factor disentanglement [ 4 , 7 , 23 , 34 , 38 , 43 ] and independent component analysis ( ICA ) [ 2 , 9 , 25–27 , 31 , 32 , 42 , 51 ] . Work in said areas focus on discovering an underlying global latent space which describes the input space . In contrast , we decompose data into a set of compositional vector spaces . This enables our approach to compose multiple instances of one factor together , as well as compose factors across distinct datasets . †short for COMposable Energy neTwork Our work is also related to a large body of existing work on unsupervised object discovery [ 5 , 10 , 14– 17 , 21 , 22 , 35 , 39 , 45 , 46 , 49 ] . Such methods seek to decompose the scene into its underlying compositional objects by independently segmenting out each individual object in the scene . In contrast , our approach represents individual objects without the need of an explicit segmentation mask , enabling our approach to represent global relations between objects . Our work draws on recent work in energy based models ( EBMs ) [ 11 , 13 , 18 , 20 , 33 , 41 , 44 , 50 ] . Our underlying energy optimization procedure to generate samples is reminiscent of Langevin sampling , which is used to sample from EBMs [ 11 , 41 , 50 ] . Most similar to our work is that of [ 12 ] , which proposes composing EBMs for compositional visual generation . Different from [ 12 ] , we study how we may discover factors in an unsupervised manner from data . 3 Composable Energy Networks . Let D = { X } be a set of images x ∈ RD . Our goal is to obtain a set Z of components { z0 , z1 , . . . , zk } for each image x representing the underlying factors of variation in the image , where each component z ∈ RM . We first discuss how we represent Z as a set of energy functions in Section 3.1 . We then discuss how COMET learns to decompose data in an unsupervised manner into a set of components in Section 3.2 . 3.1 Composing Factors of Variation as Energy Functions . In prior work for decomposing images into a set of components Z , approaches such as β-VAE [ 23 ] utilize a parameterized , non-adaptive , feedforward decoder to map the set of components to an image . Due to this , the method is unable to compose multiple instances of the same component or a larger number of components then seen during training . On the other hand , approaches such as MONET [ 5 ] utilize a decoder with shared weights to process each individual component . Unfortunately , the components are decoded independently , preventing relationships between individual components from being captured . To flexibly and generically compose a set of components Z , we require an approach that ( 1 ) utilizes a decoder that is shared across each component , such that variable-sized sets of components may be encoded , while also ( 2 ) decodes components jointly such that the relationships between individual components are covered . To construct an approach with the aforementioned desiderata , we propose to represent each component z as an energy function . We illustrate our approach and its differences from prior work in Figure 2 . Representing Components as Energy Functions . Given a single component z , we encode the factor using an energy function Eθ ( x , z ) : RD × RM → R , which is learned to assign low energy to all images x which contain the component and high energy to all other images . To obtain an image x′ with a component z , we solve for the expression x′ = argminxEθ ( x ; zi ) . Note that in contrast to prior work composing energy functions [ 12 ] , our energy functions have no probabilistic interpretation . Composing Energy Functions . Given multiple components , we represent a set of components { zi } by performing a summation ∑ iEθ ( x , zi ) over each component ’ s energy function . To obtain an image x′ for the set of components Z , we solve for the expression x′ = argmin x ∑ iEθ ( x ; zi ) . While both COMET and MONET rely on summation as a tool for composing components , MONET sums components in the image domain , while COMET sums the cost functions representing each component together . By summing cost functions , each component may combine with other components in a non-linear manner , enabling us to model more complex relationships between individual components . As a result , our generation x′ is the byproduct of jointly minimizing each individual energy function Eθ ( x , zi ) , and thus our generation contains each component . In addition , each individual component zi is parameterized by the same energy function Eθ , enabling us to model a different number of components . 3.2 Unsupervised Decomposition of Composable Energy Functions . We next discuss how COMET discovers a set of composable energy functions from an input image xi . In Section 3.1 , we discuss that a set of components Z = { zi } may be encoded as a set of composable energy functions using the expression argminx ∑ iEθ ( x ; zi ) . To discover the set of components Z in an unsupervised manner , we train by recomposing the input image xi LMSE ( θ ) = ‖ argmin x ( ∑ k Eθ ( x ; Encθ ( xi ) [ k ] ) ) − xi‖2 . ( 1 ) We utilize a learned neural network encoder Encθ ( xi ) to infer a set of components zk . In practice , evaluating the expression argminx ∑ k Eθ ( x ; Encθ ( xi ) [ k ] ) , is computationally intractable . We thus approximate the argmin operation with respect to x via N steps of gradient optimization , with an approximate optimum xNi obtained as xNi = x N−1 i − λ∇x ∑ k Eθ ( x N−1 i ; Encθ ( xi ) [ k ] ) . ( 2 ) Algorithm 1 Training algorithm for COMET . Input : data dist . pD ( x ) , step size λ , number of gradient steps N , encoder Encθ , energy function Eθ , energy components K while not converged do xi ∼ pD . Encode components zki from xi z1i , . . . , z K i ← Encθ ( xi ) . Optimize sample x0i via gradient descent : x0i ∼ U ( 0 , 1 ) for gradient step n = 1 to N do xni ← xn−1i − λ∇x ∑K k=1Eθ ( x n−1 i ; z k i ) end for . Optimize objective LMSE wrt θ : ∆θ ← ∇θ ∑N n=1 ‖x n i − xi‖2 Update θ based on ∆θ using optimizer end while In the above expression , we initialize optimization of x0i with uniform noise with λ representing the step size for each gradient step . Sample xni denotes the result after n steps of gradient descent . We train our energy functionEθ using the modified objective LMSE = ∑N n=1 ‖xni − xi‖2 , where we minimizing MSE w.r.t xni . To train parameters of Eθ we use automatic differentiation to compute gradients with respect to each optimization step depicted in Equation 2 , where for training stability , we truncate gradient backpropogation to the previous time step . We provide pseudocode for training our model in Algorithm 1 . While the approach is simple , we find that it performs favorably in both global disentanglement as well as object-level disentanglement , and requires no additional objectives or priors to shape underlying latents . We utilize the same energy architecture Eθ throughout all experiments , and present details in the appendix . Controlling Local and Global Decompositions . In many scenes , both global and local factor decompositions are valid . To control the decomposition obtained by COMET , we bias the system towards inferring local factor decompositions by utilizing i ) low latent dimensionality and ii ) positional embeddings [ 37 ] , both of which have been used in previous object discovery works [ 5 ] . This bias serves to encourage and enable models to focus on local patches of an image . We provide analysis of the effect of each inductive bias on the decomposition in Section 5.2 .

Summary: The paper proposes a novel unsupervised learning method called COMET that discovers global and local factors of variation in image datasets by representing these factors as energy functions. The experiments show that the proposed approach performs better than beta-VAE in terms of disentanglement evaluation. Qualitative examples on various datasets show that the proposed approach captures global and local object-based factors of variation and is able to generate images from various combinations of such factors well.