Unsupervised Learning of Neurosymbolic Encoders

1 INTRODUCTION . Advances in unsupervised learning have enabled the discovery of latent structures in data from a variety of domains , such as image data ( Dupont , 2018 ) , sound recordings ( Calhoun et al. , 2019 ) , and tracking data ( Luxem et al. , 2020 ) . For instance , a common approach is to use encoder-decoder frameworks , such as variational autoencoders ( VAE ) ( Kingma & Welling , 2014 ) , to identify a lowdimensional latent representation from the raw data that could contain disentangled factors of variation ( Dupont , 2018 ) or semantically meaningful clusters ( Luxem et al. , 2020 ) . Such approaches typically employ complex mappings based on neural networks , which can make it difficult to explain how the model assigns inputs to latent representations ( Zhang et al. , 2020 ) . To address this issue , we introduce unsupervised neurosymbolic representation learning , where the goal is to find a programmatically interpretable representation ( as part of a larger neurosymbolic representation ) of the raw data . We consider programs to be differentiable , symbolic models instantiated using a domain-specific language ( DSL ) , and use neurosymbolic to refer to blendings of neural and symbolic . Neurosymbolic encoders can offer a few key benefits . First , since the DSL reflects structured domain knowledge , they can often be human-interpretable ( Verma et al. , 2018 ; Shah et al. , 2020 ) . Second , by leveraging the inductive bias of the DSL , neurosymbolic encoders can potentially offer more factorized or well-separated representations of the raw data ( i.e. , the representations are more semantically meaningful ) , which has been observed in studies that used hand-crafted programmatic encoders ( Zhan et al. , 2020 ) . Our learning algorithm is grounded in the VAE framework ( Kingma & Welling , 2014 ; Mnih & Gregor , 2014 ) , where the goal is to learn a neurosymbolic encoder coupled with a standard neural decoder . A key challenge is that the space of programs is combinatorial , which we tackle via a tight integration between standard VAE training with modern program synthesis methods ( Shah et al. , 2020 ) . We further show how to incorporate ideas from adversarial information factorization ( Creswell et al. , 2017 ) and enforcing capacity constraints ( Burgess et al. , 2017 ; Dupont , 2018 ) in order to mitigate issues such as posterior and index collapse in the learned representation . We evaluate our neurosymbolic encoding approach on multiple behavior analysis domains , where the data are from challenging real-world settings and cluster interpretability is important for domain experts . Our contributions are : • We propose a novel unsupervised approach to train neurosymbolic encoders , to result in a programmatically interpretable representation of data ( as part of a neurosymbolic representation ) . • We show that our approach can significantly outperform purely neural encoders in extracting semantically meaningful representations of behavior , as measured by standard unsupervised metrics . • We further explore the flexibility of our approach , by showing that performance can be robust across different DSL designs by domain experts . • We showcase the practicality of our approach on downstream tasks , by incorporating our approach into a state-of-the-art self-supervised learning approach for behavior analysis ( Sun et al. , 2021b ) . 2 BACKGROUND . 2.1 VARIATIONAL AUTOENCODERS . We build on VAEs ( Kingma & Welling , 2014 ; Mnih & Gregor , 2014 ) , a latent variable modeling framework shown to learn effective latent representations ( also called encodings/embeddings ) ( Higgins et al. , 2016 ; Zhao et al. , 2017 ; Yingzhen & Mandt , 2018 ) and can capture the generative process ( Oord et al. , 2017 ; Vahdat & Kautz , 2020 ; Zhan et al. , 2020 ) . VAEs introduce a latent variable z , an encoder q , a decoder p✓ , and a prior distribution p on z. and ✓ are the parameters of the q and p respectively , often instantiated with neural networks . The learning objective is to maximize the evidence lower bound ( ELBO ) of the data log-likelihood : ELBO : = Eq ( z|x ) ⇥ log p✓ ( x|z ) ⇤ DKL q ( z|x ) ||p ( z ) log p ( x ) . ( 1 ) The first term in Eq . 1 is the log-density assigned to the data , while the second term is the KLdivergence between the prior and approximate posterior of z . Latent representations z are often continuous and modeled with a Gaussian prior , but z can be modeled to contain discrete dimensions as well ( Kingma et al. , 2014 ; Hu et al. , 2017 ; Dupont , 2018 ) . Our experiments are focused on behavioral tracking data in the form of trajectories , and so in practice we utilize a trajectory variant of VAEs ( Co-Reyes et al. , 2018 ; Zhan et al. , 2020 ; Sun et al. , 2021b ) , described in Section 3.4 . One challenge with VAEs ( and deep encoder-decoder models in general ) is that while the model is expressive , it is often difficult to interpret what is encoded in the latent representation z . Common approaches include taking traversals in the latent space and visualizing the resulting generations ( Burgess et al. , 2017 ) , or post-processing the latent variables using techniques such as clustering ( Luxem et al. , 2020 ) . Such techniques are post-hoc and thus can not guide ( in an interpretable way ) the encoder to be biased towards a family of structures . Some recent work have studied how to impose structure in the form of graphical models or dynamics in the latent space ( Johnson et al. , 2016 ; Deng et al. , 2017 ) , and our work can be thought of as a first step towards imposing structure in the form of symbolic knowledge encoded in a domain specific programming language . 2.2 SYNTHESIS OF DIFFERENTIABLE PROGRAMS . Our approach utilizes recent work on the synthesis of differentiable programs ( Shah et al. , 2020 ; Valkov et al. , 2018 ) , where one learns both the discrete structure of the symbolic program ( analogous to the architecture of a neural network ) as well as differentiable parameters within that structure . Our formulation of this problem closely follows that of Shah et al . ( 2020 ) . We use a domain-specific functional programming language ( DSL ) , generated with a context-free grammar ( see Figure 2 for an example ) . Programs are represented as a pair ( ↵ , ) , where ↵ is a discrete program architecture and are its real-valued parameters . We denote P as the space of symbolic programs ( i.e . programs with complete architectures ) . The semantics of a program ( ↵ , ) are given by a function [ [ ↵ ] ] ( x , ) , which is guaranteed by the semantics of the DSL to be differentiable in both x and . Like Shah et al . ( 2020 ) , we pose the problem of learning differentiable programs as search through a directed program graph G. The graph G models the top-down construction of program architectures ↵ through the repeated firing of rules of the DSL grammar , starting with an empty architecture ( represented by the “ start ” nonterminal of the grammar ) . The leaf nodes of G represent programs with complete architectures ( no nonterminals ) . Thus , P is the set of programs in the leaf nodes of G. The other nodes in G contain programs with partial architectures ( has at least one nonterminal ) . We interpret a program in a non-leaf node as being neurosymbolic , by viewing its nonterminals as representing neural networks with free parameters . The root node in G is the empty architecture ↵0 , interpreted as a fully neural program . An edge ( ↵ , ↵0 ) exists in G if one can obtain ↵0 from ↵ by applying a rule in the DSL that replaces a nonterminal in ↵ . Program synthesis in this problem setting equates to searching through G to find the optimal complete program architecture , and then learning corresponding parameters , i.e. , to find the optimal ( ↵ , ) that minimizes a combination of standard training loss ( e.g. , classification error ) and structural loss ( preferring “ simpler ” ↵ ’ s ) . Shah et al . ( 2020 ) evaluate multiple strategies for solving this problem and finds informed search using admissible neural heuristics to be the most efficient strategy ( see appendix ) . Consequently , we adopt this algorithm for our program synthesis task . 3 NEUROSYMBOLIC ENCODERS . The structure of our neurosymbolic encoder is shown in the right diagram of Figure 1 . The latent representation z = [ z , z ( ↵ , ) ] is partitioned into neurally encoded z and programmatically encoded z ( ↵ , ) . This approach boasts several advantages : • The symbolic component of the latent representation is programmatically interpretable . • The neural component can encode any residual information not captured by the program , which maintains the model ’ s capacity compared to standard deep encoders . • By incorporating a modular design , we can leverage state-of-the-art learning algorithms for both differentiable encoder-decoder training and program synthesis . We denote q and q ( ↵ , ) as the neural and symbolic encoders respectively ( see Figure 1 ) , where z ⇠ q ( ·|x ) and z ( ↵ , ) ⇠ q ( ↵ , ) ( ·|x ) . q is instantiated with a neural network , but q ( ↵ , ) is a differentiable program with architecture ↵ and parameters in some program space P defined by a DSL . Given an unlabeled training set of x ’ s , the VAE learning objective in Eq . 1 then becomes : max , ( ↵ , ) , ✓ E q ( z |x ) q ( ↵ , ) ( z ( ↵ , ) |x ) ⇥ log p✓ ( x|z , z ( ↵ , ) ) | { z } reconstruction loss ⇤ DKL q ( z |x ) ||p ( z ) | { z } regularization for neural latent DKL q ( ↵ , ) ( z ( ↵ , ) |x ) ||p ( z ( ↵ , ) ) | { z } regularization for symbolic latent . ( 2 ) Compared to the standard VAE objective in Eq . 1 for a single neural encoder , Eq . 2 has separate KL-divergence terms for the neural and programmatic encoders . 3.1 LEARNING ALGORITHM . The challenge with solving for Eq . 2 is that while ( , , ✓ ) can be optimized via back-propagation with ↵ fixed , optimizing for ↵ is a discrete optimization problem . Since it is difficult to jointly optimize over both continuous and discrete spaces , we take an iterative , alternating optimization approach . We start with a fully neural program ( one with empty architecture ↵0 ) trained using standard differentiable optimization ( Figure 1 , Step 1 ) . We then gradually make it more symbolic ( Figure 1 , Step 2 ) by finding a program that is a child of the current program in G ( more symbolic by construction of G ) that outputs as similar to the current latent representations as possible : min ↵0 : ( ↵ , ↵0 ) 2G , 0 Lsupervised q ( ↵ , ) ( x ) , q ( ↵0 , 0 ) ( x ) , ( 3 ) which can be viewed as a form of distillation ( from less symbolic to more symbolic programs ) via matching the input/output behavior . We solve Eq . 3 by enumerating over all child programs and selecting the best one , which is similar to iteratively-deepened depth-first search in Shah et al . ( 2020 ) ( see appendix ) . We alternate between optimizing Eq . 2 and Eq . 3 until we obtain a complete program . Algorithm 1 outlines this procedure and is guaranteed to terminate if G is finite by specifying a maximum program depth . We chose this optimization procedure for two reasons . First , it maximally leverages state-of-the-art tools in both differentiable latent variable modeling ( VAE-style training ) and supervised program synthesis , leading to tractable algorithm design . Second , this procedure never makes a drastic change to the program architecture , leading to relatively stable learning behavior across iterations . Algorithm 1 Learning a neurosymbolic encoder 1 : Input : program space P , program graph G 2 : initialize , , ✓ , ↵ = ↵0 ( empty architecture ) 3 : while ↵ is not complete do 4 : , , ✓ optimize Eq . 2 with ↵ fixed 5 : ( ↵ , ) optimize Eq . 3 6 : end while 7 : , , ✓ optimize Eq . 2 with complete ↵ 8 : Return : encoder { q , q ( ↵ , ) } Algorithm 2 Learning a neurosymbolic encoder with k programs 1 : Input : program space P , program graph G , k 2 : for i = 1 .. k do 3 : fix programs { q ( ↵1 , 1 ) , . . . , q ( ↵i 1 , i 1 ) } 4 : execute Algorithm 1 to learn q ( ↵i , i ) 5 : remove q ( ↵i , i ) from P to avoid redundancies 6 : end for 7 : Return : encoder { q , q ( ↵1 , 1 ) , . . . , q ( ↵k , k ) }

Summary: The authors propose an unsupervised approach to train neurosymbolic encoders to obtain a programmatically interpretable representation using the dictionary of a domain-speciﬁc language (DSL). The experimental results show that the proposed method can outperform baseline neural encoders in extracting semantically meaningful representations of behavior in CalMS21 (mice) dataset. The results also show that the performance can be robust across different DSL designs by domain experts.