Title: LEARNING TO CONTROL LATENT REPRESENTATIONS FOR FEW-SHOT LEARNING OF NAMED ENTITIES

Abstract: 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 := (K mem_size×key_dim , V mem_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] = 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)

Figure 1 : Architecture of our approach based on a neural encoder that generates input representations and a trainable controller that learns to store and maintain them in an external memory formed by key-value pairs. The memory is modified via operations a, each of which returns a reward signal r.

neural network that transforms the raw input x into the last hidden state of the following recurrence h i = LST M (x, h i-1 ), where s = h i 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) = Sof tmax(s • K i ) i π θ ∼ argmax i (π θ (i|s)).

Then, the controller executes the memory operation a(i π θ , s, y) at position i π θ considering the Definition 2.4. During inference, ŷ = 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 (K t-1 ) to an updated state (K t ). More formaly, our reward function r t (s, a) corresponds to the reduction of entropy H when softly attending over the memory entries before and after executing the operation a.

h Kt = s • K t att Kt = Sof tmax(h Kt )

We use the dot-product attention score Luong, Pham, and Manning (2015) to compute the logits h Kt over the memory keys and use Softmax to generate the probability distribution att Kt . 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 .

r t = -H(att Kt ) -H(att Kt-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 π θ , 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(θ) = E π θ T t=1 r t (s, a) .

∇ θ J(θ) ≈ 1 N N n=1 T t=1 ∇ θ log π θ (i t n |s t n ) T t =t r t (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.

5. RELATED WORK:

The idea of training a controller that interacts with a memory to incrementally allocate distributional representations builds upon a wide range of research in machine learning and memory augmented algorithms. Kraska et al. (2018) introduced the concept of learned indices structures as a learning architecture that predicts the position or existence of records. Their main goal is to prevent too many collisions from being mapped to the same position inside the hash index. This is because collisions often cause memory overhead when traversing a linked list or require the allocation of additional memory for storing more records. They approach this problem as supervised learning of the cumulative distribution function of hash keys, which leads to minimizing the number of collisions. In contrast, we propose the use of an associative array as an indexed structure that supports generalization, so the collision of multiple observations to a similar hash key if adequate for predicting the class of similar input data. While the idea of learning hash functions as neural networks is not new, existing work mainly focused on learning a better hash-function to map observations into low-dimensional embeddings for similarity search assuming a fixed data distribution Qian et al. (2014) 

6. CONCLUSIONS:

Learning new tasks without forgetting the previously ones references the plasticity of our own brain to retain previously acquired knowledge. In this work, we show a learning system that stores information about named entities with two levels of representation: dense and sparse. While a dense memory captures features that are shared across tasks, sparse memory incorporates embedding vectors and prevents them from being overwritten by global updates during backpropagation. To do this, we rely on the sparse structure of associative arrays to locally update a small set of memory entries and on the use of a trainable controller that manages the transition of the sparse memory to a state of low entropy that also addresses the learning task. This enables a model to incrementally learn in the presence of a few observations per class. We experimentally show that our system obtains accurate results for few-shot learning of entities in the Stanford Task-Oriented Dialogue dataset in comparison with state-of-the-art few shot learning algorithms (Memory Augmented Neural Networks and Matching Networks).

:

https://nlp.stanford.edu/software/CRF-NER.html

Figures and Tables:

Average accuracy for few-shot learning in the STDO dataset.

Table content: <table><tr><td>Model</td><td colspan="4">5-way 1-shot 5-way 5-shot 20-way 1-shot 20-way 5-shot</td></tr><tr><td>Memory Augmented NN</td><td>63.1%</td><td>66.7%</td><td>57.2%</td><td>61.3%</td></tr><tr><td>Matching Networks</td><td>67.3%</td><td>71.4%</td><td>62.2%</td><td>68.3%</td></tr><tr><td>Learning to Control</td><td>71.5%</td><td>77.1%</td><td>65.3%</td><td>75.8%</td></tr></table>

. To our knowledge, it has not been explored yet if it is possible to learn a hash index according to a data distribution in which few training observations are presented per class during training.Deep Neural Networks are models that solve classification problems with non-linearly separable classes. These are shown to be good for problems such as object detectionKrizhevsky, Sutskever, and Hinton (2017), language translationJohnson et al. (2016) and image generationGoodfellow et al. (2014). Recently, recurrent neural networks (RNN) have been studied to manage the states between time steps using an attention mechanism that addresses similar contentBahdanau, Cho, and Bengio (2014). Some similar approaches have also studied the problem of few-shot learningSantoro et al. (2016)Koch, Zemel, and Salakhutdinov (2015)Vinyals et al. (2016), a representative example is Matching NetworksVinyals et al. (2016), which uses an attention mechanism to augment neural networks for set-to-set learning.A relevant research topic to this work is the idea of memory augmented networks that extend its capabilities by coupling an external memory. For example,Kaiser et al. (2017) also proposes a key-value memory, but with no controller mechanism for adequate training of few-shot learning tasks. Also, Neural Turing Machines (NTM)Graves, Wayne, and Danihelka (2014) are differentiable architectures allowing efficient training with gradient descend and showing important properties for associative recall for learning different sequential patterns. Although they have important properties for one-shot learning given their sequential memory management, the supervised nature of its training still shows issues related to catastrophic forgetting.

Table content: <table/>

Peer review (iclr2020):
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]. From a high-level perspective, memory networks have played an important role in building dialogue managers (and in several applications) and their proposed expansion over [Kaiser, et al., ICLR17] in the context of few-shot, sparse memory architectures seems sensible. Additionally, the empirical results o NER seem promising. However, in my opinion, this paper is making contributions along two dimensions, but neither of them convincingly as detailed below. First, one potential point of emphasis is the ‘learning to control’ mechanism through PG RL being better than [Kaiser, et al., ICLR17] using with a LRU-like policy. To do this, it would make sense to both compare to the datasets in the previous work (Omniglot, WMT — even if it does require additional architectures that are pretty straightforward expansions) and to include this architecture in the submitted paper as there is a publicly available dataset (https://github.com/himani-arora/learning_to_remember_rare_events). While the synthetic dataset probably isn’t necessary, one could easily think of synthetic datasets that would clearly contrast with the existing method and conduct ablation/configuration studies regarding different number of neighbors, size of memory, exact vs. approximate nearest neighbor for lookup, etc. A second point of interest is the NER task. From this perspective, there is some existing work (e.g., [Fritzler, Logacheva & Kretov, SAC19; Hou, et al., https://arxiv.org/pdf/1906.08711.pdf (Not required as not accepted yet); Hofer, et al., https://arxiv.org/pdf/1811.05468.pdf (Not clear if accepted)]. While these are largely preliminary works, they do provide datasets (and there are many others) for NER that would potentially make a convincing case for this being a state-of-the-art approach for few-shot learning in NER. However, I am not sure that being the state-of-the-art for few-shot NER would be a sufficient contribution for ICLR, even if done convincingly. If there was something in the algorithm specific to sequence prediction, then I think few-shot sequence prediction would be a sufficient contribution, but given the current architecture, I would think that the bar is a methodological contribution to adding sparse memory units (via RL in this case) is the acceptable level of contribution. Finally, my belief that the contribution is too narrow and not sufficiently developed is supported by the writing seeming rushed. The paper makes much more sense after reading [Kaiser, et al., ICLR17] as Sections 2, 3 of the submission are missing the aspects of a clear constrast to [Kaiser, et al., ICLR17], a working example to provide readers a more intuitive understanding, a clear example showing dimensionality and notation (beyond what is in Figure 1). There are also many typos (e.g., “mammals exel”, “trainale controller”, “Rathern then”, amongst others) and general need for proofreading. However, this is not the determining factor in the paper. Basically, I think the contribution is insufficient in terms of scope and convincingness of the contribution. Thus, I recommend rejecting in its current form — even if the underlying idea is promising and worth developing further.