Learning to Control Latent Representations for Few-Shot Learning of Named Entities

Humans excel in continuously learning with small data without forgetting how to solve old problems . However , neural networks require large datasets to compute latent representations across different tasks while minimizing a loss function . For example , a natural language understanding ( NLU ) system will often deal with emerging entities during its deployment as interactions with users in realistic scenarios will generate new and infrequent names , events , and locations . Here , we address this scenario by introducing a RL trainable controller that disentangles the representation learning of a neural encoder from its memory management role . Our proposed solution is straightforward and simple : we train a controller to execute an optimal sequence of read and write operations on an external memory with the goal of leveraging diverse activations from the past and provide accurate predictions . Our approach is named Learning to Control ( LTC ) and allows few-shot learning with two degrees of memory plasticity . We experimentally show that our system obtains accurate results for few-shot learning of entity recognition in the Stanford Task-Oriented Dialogue dataset . 1 MOTIVATION . Today , supervised models have problems incorporating new tasks over time while protecting previously acquired knowledge . This is because these algorithms require that all data is given prior to training . This becomes a problem in the presence of more general scenarios in which new classes emerge during training or data exhibit long-tailed distributions . Hence , classes with large support dominate the learning of gradient-based representations causing the catastrophic forgetting of under-represented classes , Kirkpatrick et al . ( 2016 ) . Unlike deep neural networks , humans and other mammals exel in learn incrementally with small data . Biological evidence suggests that the process of acquiring new skills occurs in different brain areas with at least two degrees of plasticity , Wixted et al . ( 2018 ) : 1 ) A slow memory that is dense and requires extensive practive . 2 ) A fast memory that is sparse and stores volatile information . Inspired by these observations , we introduce a trainable controller , Learning to Control ( LTC ) , that learns to interact with both a slow memory ( consisting of neural encoders ) and a fast memory ( consisting of an associative array storing key-value pairs ) . We experimentally show the advantage of using LTC for a few-shot learning setup . Figure 1 depicts the network architecture of our model . We present the following contributions : • We propose a novel architecture that uses a trainale controller to manipulate latent representations in an external memory ( Section 3 ) . • We introduce the use of a reward signal that is proportional to the average reduction of entropy when attending the memory entries . This enables us to propose a reinforcement learning approach that learns a policy based on interactions with the external memory . • We show the generality of our solution for the few-shot learning of entities in the Stanford Task-Oriented Dialogue dataset . 2 DENSE AND SPARSE MEMORIES . The proposed architecture considers the use of two types of memories that behave differently during backpropagation . First , we define a memory as follows . Definition 2.1 . ( Memory ) Given a neural model f ( x , y , θ ) that maps an observation x to the learning task y and is parameterized by θ . Then , the memory of f ( x , y , θ ) is a collection of co-activations responding similarly to the reocurring patterns associated to the learning task y . A Neural Network uses most of its memory to remember patterns in the dataset . However , our approach distinguishes between dense and sparse memories . Definition 2.2 . ( Dense Memory ) A dense memory is a type of memory , in which all its learnable parameters θ have the possibility of being updated during training . We want to preserve latent representations of the input between training steps but a dense memory will likely overwrite them based on frequent global updates . We then need a sparse memory with addressable memory entries and sparse updates . Definition 2.3 . ( Sparse Memory ) Given a model f ( x , y , θ ) , let a ( x , y ) be a memory operation which returns k updating parameters to incorporate the pair ( x , y ) . Then , a sparse memory is a type of memory that satisfies the sparse updating constraint : 0 < |a ( x , y ) | |θ| where |θ| is the total number of trainable parameters . We use the sparse memory introduced in Kaiser et al . ( 2017 ) which consists of two arrays of size mem_size for storing keys ( K ) and values ( V ) . Each key has a dimensionality of key_dim and each value is a scalar representing a class label . In our experiments , we fixed the number of updating parameters , i.e . keys , to be very small ( k = 10 ) in comparison to mem_size in order to satisfy the sparse updating constraint , k mem_size . The final form of the sparse memory is as follows . M : = ( Kmem_size×key_dim , Vmem_size ) The state of a sparse memory changes based on the following memory operations . Definition 2.4 . ( Memory Operations ) Let s be a input embedding that we want to incorporate into M , y be a class label , and i be the memory index of the most similar key to s based on the Cosine similarity . 1 . If V [ i ] 6= y , the write ( i , s , y ) operation registers a new class with M [ i ] = 〈s , y〉 . 2 . If V [ i ] = y , the update ( i , s , y ) operation modifies the stored key-based representation K [ i ] with M [ i ] = 〈‖K [ i ] + s‖ , y〉 . 3 LEARNING TO CONTROL . The goal of Learning To Control ( LTC ) is to learn a strategy to manage the entries of a sparse memory . We characterize this decision process with a memory policy πθ which represents the probability of running an memory operation given the current state of the memory . Rathern than following a heuristics ( e.g. , overwritting in the oldest memory entry as in Kaiser et al . ( 2017 ) ) , learning such policy has the advantage of optimizing the allocation of infrequent classes considering the limited number of memory entries in M . Indeed , πθ transitions M to a new configuration in which its keys are more adapted to store and maintain informative latent vectors . To clarify the connection between the components , let us consider the forward and backward computations involved in LTC . 3.1 FORWARD COMPUTATION . As illustrated in Figure 1 , information flows through the hidden layers of the neural encoder , feeds the controller , and reaches the external memory M . The encoder is a Long-Short Term Memory ( LSTM ) neural network that transforms the raw input x into the last hidden state of the following recurrence hi = LSTM ( x , hi−1 ) , where s = hi is the state of the environment prior to any memory operation and corresponds to the last hidden state of the LSTM encoder . Which memory entries should the controller update to incorporate the learning task ( s , y ) ? Our approach is to induce the probablity distribution πθ ( i|s ) over the memory entries given the state s and sample from there the most likely memory index iπθ πθ ( i|s ) = Softmax ( s ·Ki ) iπθ ∼ argmaxi ( πθ ( i|s ) ) . Then , the controller executes the memory operation a ( iπθ , s , y ) at position iπθ considering the Definition 2.4 . During inference , ŷ = V [ iπθ ] returns the predicted class of the raw input x . 3.2 BACKWARD COMPUTATION . The forward step performs an operation a ( iπθ , s , y ) in the sparse memory and modifies its keys to integrate latent representations . The backward step estimates the objective function J ( θ ) and computes its gradients with respect to the trainable parameters of the encoder and controller modules . In the case of the controller , we want to compute the gradient vector ∇θJ ( θ ) that updates the stochastic policy πθ in a direction that maximizes J ( θ ) during few-shot learning . For our purposes , each memory operation a returns a reward r that measures the efficacy of this operation to reduce uncertainty about which of the keys in K will lead the transition between memory states . This means going from a state in which no operation took place in memory ( Kt−1 ) to an updated state ( Kt ) . More formaly , our reward function rt ( s , a ) corresponds to the reduction of entropy H when softly attending over the memory entries before and after executing the operation a. hKt = s ·Kt attKt = Softmax ( hKt ) We use the dot-product attention score Luong , Pham , and Manning ( 2015 ) to compute the logits hKt over the memory keys and use Softmax to generate the probability distribution attKt . Zhu et al . ( 2005 ) demonstrates that models that show high entropy in their output distributions often fail to discriminate classes because of exhibiting uniform-like distributions and a high degree of uncertainty . Thus , we use the following entropy-based reward signal to encourage the use of few memory entries to modify the sparse memory M . rt = − ( H ( attKt ) −H ( attKt−1 ) ) Based on memory operations , we can approximate the objective function from a batch of sampled episodes of length T as the expected value of the cumulative sum of rewards under the policy πθ , J ( θ ) = Eπθ [ T∑ t=1 rt ( s , a ) ] . The REINFORCE algorithm Williams ( 1992 ) allow us to directly take the gradients of the objective function J ( θ ) as described in Equation 1 and used them to perform backpropagation . ∇θJ ( θ ) ≈ 1 N N∑ n=1 ( T∑ t=1 ∇θ log πθ ( itn|stn ) ) ( T∑ t′=t rt′ ( s t′ n , a t′ n ) ) ( 1 ) 4 EXPERIMENTS . We evaluate our proposed method on the Stanford Task-Oriented Dialogue Dataset ( STDO ) Eric et al . ( 2017 ) which consists of 3 , 031 dialogues in the domain of an in-car assistant that provides automatic responses to the requests of a driver considering weather , point-of-interest , and scheduled domains . 4.1 BASELINES . We compare Learning to Control ( LTC ) with the following baselines : 1 ) Matching Networks ( MN ) , Vinyals et al . ( 2016 ) : This algorithm provides one-shot learning capabilities by jointly mapping a small labeled support set and an unlabeled example to its most likely label , and 2 ) Memory Augmented Neural Networks ( MANN ) , Santoro et al . ( 2016 ) : This algorithm is designed to provide one-shot learning based on a Neural Turing Machine and a curriculum training regime . 4.2 FEW-SHOT LEARNING FOR ENTITY RECOGNITION . We study the problem of sequence labeling to find the best label sequence ( named entities ) for a given input sentence considering an LSTM encoder augmented with an external memory to store hidden representations for each recurrent unit . We use the Stanford Named Entity Recognizer ( NER ) 1 to augment the STDO dataset with 7 classes ( Location , Person , Organization , Money , Percent , Date , Time ) and an non-entity class . This results in 15 , 928 observations , 8 classes , and an average sequence length of 44 words . We change the format of the STOD dataset to simulate a scenario in which classes obtain incremental support during training . This is done by forming episodes of labeled examples that show a uniform distribution over k classes ( e.g. , 5-way ) , so we can track the number of times a particular class was presented to the model during training ( e.g. , 1-shot learning ) . This format was proposed by Santoro et al . ( 2016 ) to study few-shot learning . The goal of this experiment is to learn with small data , so we study the learning behavior of LTC with an external memory of 1 , 000 entries and in relation to two models designed to provide few-shot learning capabilities : Memory Augmented Neural Networks , Santoro et al . ( 2016 ) and Matching Networks , Vinyals et al . ( 2016 ) . Table 1 shows that LTC provides an average improvement of +5.2 in the STOD dataset with respect to the MN neural model , which is consistently the second most accurate option for both 5-way and 20-way classification and when the number of classes ranges from small ( 5-way ) to large ( 20-way ) . Similarly , LTC also shows an average improvement of +10.35 % with respect to the MANN model . For the scenario of learning with infrequent classes , LTC outperforms MN when only 1 training observation is presented ( 1-shot ) and 5 and 20 classes are available during training by +4.2 % and +3.1 % , respectively . The performance advantage in the results can be explained by the extra memory capacity of LTC and represent an opportunity to explore later what is the minimum memory size to continuously learn with small data , an even more challenging problem to study in the future . 1https : //nlp.stanford.edu/software/CRF-NER.html

Summary: The authors build on work regarding few-shot learning with memory augmented networks, specifically [Kaiser, et al., ICLR17] where the goal is to learn a memory address mapping such that generalization is achieved by finding the nearest neighbor memory address when predicting the label. For correct predictions, the memory key is updated to include the associated (predicted) address while new memory locations are written for mistakes. Whereas [Kaiser, et al., ICLR17] follows a LRU-like procedure for replacing memory, the current work proposes performing policy-gradient RL where the action space is the memory locations and the reward is reduction of entropy over the memory address assignment distribution over the memory locations. This approach is empirically studied for an RNN approach to NER, specifically considering few-shot learning for NER in the Stanford Task-Oriented Dialogue (STOD) dataset — showing non-negligible improvements over Memory Augmented Networks [Santoro, et al., ICML16] and Matching Networks [Vinyals, et al., NeurIPS16].