Can a Fruit Fly Learn Word Embeddings?

1 INTRODUCTION . Deep learning has made tremendous advances in computer vision , natural language processing and many other areas . While taking high-level inspiration from biology , the current generation of deep learning methods are not necessarily biologically realistic . This raises the question whether biological systems can further inform the development of new network architectures and learning algorithms that can lead to competitive performance on machine learning tasks or offer additional insights into intelligent behavior . Our work is inspired by this motivation . We study a well-established neurobiological network motif from the fruit fly brain and investigate the possibility of reusing it for solving common machine learning tasks in NLP . We consider this exercise as a toy model example illustrating the possibility of “ reprogramming ” of naturally occurring algorithms and behaviors ( clustering combinations of input stimuli from olfaction , vision , and thermo-hydro sensory system ) into a target algorithm of interest ( learning word embeddings from raw text ) that the original biological organism does not naturally engage in . The mushroom body ( MB ) is a major area of the brain responsible for processing of sensory information in fruit flies . It receives inputs from a set of projection neurons ( PN ) conveying information ∗Yuchen Liang is an AI Horizons Scholar , part of the Rensselaer-IBM AI Research Collaboration ( AIRC ) . from several sensory modalities . The major modality is olfaction [ 2 ] , but there are also inputs from the PN responsible for sensing temperature and humidity [ 29 ] , as well as visual inputs [ 45 ; 6 ] . These sensory inputs are forwarded to a population of approximately 2000 Kenyon cells ( KCs ) through a set of synaptic weights [ 26 ] . KCs are reciprocally connected through an anterior paired lateral ( APL ) neuron , which sends a strong inhibitory signal back to KCs . This recurrent network effectively implements winner-takes-all competition between KCs , and silences all but a small fraction of top activated neurons [ 8 ] . This is the network motif that we study in this paper ; its schematic is shown in Fig . 1 . KCs also send their outputs to mushroom body output neurons ( MBONs ) , but this part of the MB network is not included into our mathematical model . In computational linguistics there is a long tradition [ 19 ] of using distributional properties of linguistic units for quantifying semantic similarities between them , as summarized in the famous quote by JR Firth : “ a word is characterized by the company it keeps ” [ 14 ] . This idea has led to powerful tools such as Latent Semantic Analysis [ 9 ] , topic modelling [ 3 ] , and language models like word2vec [ 30 ] , GloVe [ 34 ] , and , more recently , BERT [ 10 ] which relies on the Transformer model [ 44 ] . Specifically word2vec models are trained to maximize the likelihood of a word given its context , GloVe models utilize global word-word co-occurence statistics , and BERT uses a deep neural network with attention to predict masked words ( and the next sentence ) . As such , all these methods utilize the correlations between individual words and their context in order to learn useful word embeddings . In our work we ask the following question : can the correlations between words and their contexts be extracted from raw text by the biological network of KCs , shown in Fig . 1 ? Further , how do the word representations learned by KCs differ from those obtained by existing NLP methods ? Although this network has evolved to process sensory stimuli from olfaction and other modalities and not to “ understand ” language it uses a general purpose algorithm to embed inputs ( from different modalities ) into a high dimensional space with several desirable properties , which we discuss below . Our approach relies on a recent proposal that the recurrent network of mutually inhibited KCs can be used as a “ biological ” model for generating sparse binary hash codes for the input data presented at the projection neuron layer [ 8 ] . It was argued that a matrix of random weights projecting from PN layer into the KCs layer leads to the highly desirable property of making the generated hash codes locality sensitive , i.e. , placing similar inputs close to each other in the embedding space and pushing distinct stimuli far apart . A subsequent study [ 39 ] has demonstrated that the locality sensitivity of the hash codes can be significantly increased , compared to the random case , if the matrix of weights from PN to KCs is learned from data . The idea of using the network of KCs with random projections for NLP tasks has also been previously explored in [ 37 ] , see discussion in section 6 . Biologically , there is an ongoing debate in the neuroscience community regarding whether these projections are random . For instance , [ 5 ] argues for the random model , while [ 47 ] presents evidence of the non-random structure of this network , which is related to the frequency of presented odors . Since the goal of our work is to build a useful AI system and not mimic every detail of the biological system , we adopt the data-driven synaptic weight strategy even if fruit flies may use random projections . As is clearly demonstrated in [ 39 ] , learned synapses lead to better performance . Our main contributions in this work are the following : 1 . Inspired by the fruit fly network , we propose an algorithm that makes it possible to generate binary ( as opposed to continuous ) word embeddings for words and their context . We systematically evaluate the performance of this algorithm on word similarity task , word-sense disambiguation , and document classification . 2 . We demonstrate that our binary embeddings result in tighter and better separated clusters of concepts compared to continuous GloVe embeddings , and stand in line with clustering properties of binarized versions of GloVe . 3 . We show that training the fruit fly network requires an order of magnitude smaller compute time than training the classical NLP architectures , like BERT , at the expense of relatively small decrease in classification accuracy . 2 LEARNING ALGORITHM . Consider a training corpus . Each sentence can be decomposed into a collection of w-grams of consecutive words . If the word tokens come from a predefined vocabulary of size Nvoc , the input to the algorithm is a vector of size 2Nvoc . This vector consists of two blocks : the context ( the first Nvoc elements ) , and the target ( the remaining Nvoc elements ) ; see Fig . 2 . In this work w is assumed to be an odd integer , and the target word is assumed to be the center of the w-gram . The target word is one-hot encoded in the target block , and the context words are binary encoded as a bag of words in the context block ( no positional information is used ) . The window w slides along the text corpus , and for each position generates a training vector vA = { vAi } 2Nvoc i=1 , where the index A enumerates different w-grams , and index i enumerates positions in the context-target vector . These training vectors are passed to the learning algorithm . The goal of the algorithm is to learn correlations between the context and the target blocks . 2.1 MATHEMATICAL FORMULATION . Mathematically , the objective of the training algorithm is to distribute a set of context-target pairs among K buckets , so that similar pairs end up in similar buckets . In order to achieve this , the learning algorithm takes two inputs : a set of training vectors vA ∈ { 0 , 1 } 2Nvoc , and a vector of occurrence probabilities p = { pi = f ( i mod Nvoc ) } 2Nvoc i=1 ∈ R2Nvoc , where fj is the probability of observing word j in the training corpus1 . The learning can be formalized as a minimization of the 1In our notation vector f has Nvoc elements , while vector p has 2Nvoc elements . Given that index i runs from 1 to 2Nvoc , notation pi = f ( i mod Nvoc ) is a mathematical way to concatenate two vectors f into a twice longer vector p. energy function , see [ 39 ] for additional details , defined by E = − ∑ A∈data 〈 Wµ̂ , v A/p 〉 〈 Wµ̂ , Wµ̂ 〉1/2 , where µ̂ = arg max µ 〈 Wµ , v A 〉 ( 1 ) In this equation W ∈ RK×2Nvoc is a matrix of synaptic connections , given as W = { Wµ } = { Wµi } , projecting from PN layer ( individual neurons in the layer are denoted by the index i ) to the KC layer ( individual neurons in the KC layer are denoted by the index µ ) . There are 2Nvoc neurons in the PN layer andK neurons in the KC layer . The inner product 〈X , Y〉 = ∑2Nvoc i=1 XiYi is defined as a contraction over index i of PN cells . In the numerator of the energy function the binary encoded w-gram is divided by the probabilities of occurrences of individual words element-wise , so that the numerator can be written as 〈 Wµ̂ , v A/p 〉 = 2Nvoc∑ i=1 Wµ̂i vAi pi Probabilities p are calculated based on the frequencies of words in the training corpus . The vocabulary contains Nvoc most frequent words in the corpus , thus all the elements of pi are non-zero and the element-wise division is well defined . Intuitively , the goal of the training algorithm is to adjust the weights of the neural network so that they are aligned with w-grams that are frequently present in the training corpus . We rely on the assumption that semantically related w-grams share several “ core ” words , while a few individual words might be substituted by synonyms/antonyms . The minimization of the energy function ( 1 ) is accomplished by the iterative update of the weights satisfying the following learning rule [ 25 ; 39 ; 17 ] ∆Wµi = εg [ ∑ j Wµjv A j ] [ vAi pi − ( ∑ j Wµj vAj pj ) Wµi ] ( 2 ) In this equation the activation function is equal to one for a maximally driven hidden unit ( Kenyon cell ) , and is equal to zero otherwise g [ xµ ] = δµ , µ̂ , where µ̂ = arg max µ [ xµ ] ( 3 ) The learning rate is denoted by ε , and δµ , µ̂ is a Kronecker delta symbol . 2.2 BIO-HASHING . After learning is complete the hash codes for the inputs can be generated in the following way . Given the binary encoded w-gram vA , Hµ = { 1 , if 〈Wµ , vA〉 in the top k of all KCs activations 0 , otherwise ( 4 ) This is a crude mathematical approximation of the biological computation performed by the PN– KC–APL neural network [ 8 ; 39 ] . An input vA generates an input current 〈Wµ , vA〉 into the KC neurons using feedforward weightsWµi . The recurrent network of KCs and the APL neuron silences all but a small fraction of KCs . Those cells that remain active are assigned state 1 , while the rest of the KCs are assigned the inactive state 0 . Notice , that equation ( 4 ) makes it possible to generate the hash codes for both individual words ( static word embeddings like word2vec and GloVe ) and phrases ( similar to Transformer models ) . In the static case , the input vA has all zeros in the context block and a one-hot encoded word in the target block . In the context-dependent case , both blocks have binary encoded input words . Importantly , both context-dependent and static embeddings are mapped into the same space of sparse binary hash codes ( a vector of K elements , with k ones in it ) . We show below that these hash codes capture semantic meaning of the target word and the context in which it is used . For the rest of the paper we refer to the parameter k in equation ( 4 ) as the hash length . In order to provide an intuition behind the learning algorithm defined by the energy function ( 1 ) and weight update rule ( 2 ) and connect it to some of the existing methods in machine learning , consider the limit when all the words have equal probabilities in the training corpus , pi = 1Nvoc . In this limit the energy function ( 1 ) reduces to the familiar spherical K-means clustering algorithm [ 11 ] . In this limit the weights of each KC correspond to the centroids of the clusters of context-target vectors . The hashing rule ( 4 ) assigns active state 1 to the k closest centroids ( and inactive state 0 to the remaining ones ) , defined with respect to cosine similarity distance . In this simple limit the learning algorithm that we use can be viewed as a biologically plausible implementation of this classical algorithm . For real datasets the probabilities of words are different , thus this correspondence does not hold . Notice that division by the probability appears only in the expression for the energy , but not in the definition of µ̂ in equation ( 1 ) . Equivalently , division by pi appears in the second bracket of equation ( 2 ) , but not in the argument of the activation function g [ xµ ] . Thus , in the general case ( for different word probabilities pi ) our algorithm is not equivalent to spherical K-means on context-target vectors rescaled by the probabilities . Rather , in the general case , the closest centroid is found for a given context-target vector ( via the definition of µ̂ in equation ( 1 ) - no pi involved ) , but the “ updates of the position ” of that centroid are computed by enhancing the contributions of rare words ( small pi ) and suppressing the contributions of frequent words ( large pi ) . Empirically , we have found that division by the probabilities improves performance of our method compared to the case of spherical K-means ( when the factor 1/p is removed from the algorithm ) .

Summary: Although the paper does not say so, my understanding is that the proposed word embedding method actually first perform Kmeans-like clustering on the context vector shown in Figure 2. In the binary word embedding of each word, we set the dimensions corresponding to the k closest cluster centers to 1 and 0 otherwise. Most parts of the paper are describing how the simple word embedding method is related to the neural system of a fruit fly and showing the method achieves comparable performance compared with GloVe in word similarity tasks, context-sensitive word similarity datasets, and document classification tasks.