Neural Speed Reading with Structural-Jump-LSTM

1 INTRODUCTION . Recurrent neural networks ( RNNs ) are a popular model for processing sequential data . The Gated Recurrent Unit ( GRU ) ( Chung et al. , 2014 ) and Long Short Term Memory ( LSTM ) ( Hochreiter & Schmidhuber , 1997 ) are RNN units developed for learning long term dependencies by reducing the problem of vanishing gradients during training . However , both GRU and LSTM incur fairly expensive computational costs , with e.g . LSTM requiring the computation of 4 fully connected layers for each input it reads , independently of the input ’ s importance for the overall task . Based on the idea that not all inputs are equally important , and that relevant information can be spread throughout the input sequence , attention mechanisms were developed ( Bahdanau et al. , 2015 ) to help the network focus on important parts of the input . With soft attention , all inputs are read , but the attention mechanism is fully differentiable . In comparison , hard attention completely ignores part of the input sequence . Hard attention mechanisms have been considered in many areas , ranging from computer vision ( Mnih et al. , 2014 ; Campos et al. , 2018 ) where the model learns what parts of the image it should focus on , to natural language processing ( NLP ) , such as text classification and question answering ( Yu et al. , 2017 ; Campos et al. , 2018 ; Yu et al. , 2018 ) , where the model learns which part of a document it can ignore . With hard attention , the RNN has fewer state updates , and therefore fewer floating point operations ( FLOPs ) are needed for inference . This is often denoted as speed reading : obtaining the same accuracy while using ( far ) fewer FLOPs ( Yu et al. , 2017 ; 2018 ; Seo et al. , 2018 ; Huang et al. , 2017 ; Fu & Ma , 2018 ) . Prior work on speed reading processes text as chunks of either individual words or blocks of contiguous words . If the chunk being read is important enough , a full state update is performed ; if not , the chunk is either ignored or a very limited amount of computations are done . This is followed by an action aiming to speed up the reading , e.g . skipping or jumping forward in text . Inspired by human speed reading , we contribute an RNN speed reading model that ignores unimportant words in important sections , while also being able to jump past unimportant sections of the text . Our model , called Structural-Jump-LSTM1 , both skips and jumps over dynamically defined chunks of text as follows : ( a ) it can skip individual words , after reading them , but before updating the RNN state ; ( b ) it uses the punctuation structure of the text to define dynamically spaced jumps to the next sub-sentence separator ( , ; ) , end of sentence symbol ( . ! ? ) , or the end of the text . An extensive experimental evaluation against all state-of-the-art speed reading models ( Seo et al. , 2018 ; Yu et al. , 2017 ; 2018 ; Fu & Ma , 2018 ; Huang et al. , 2017 ) , shows that our Structural-JumpLSTM of dynamically spaced jumps and word level skipping leads to large FLOP reductions while maintaining the same or better reading accuracy than a vanilla LSTM that reads the full text . 2 RELATED WORK . Prior work on speed reading can be broadly clustered into two groups : 1 ) jumping based models , and 2 ) word level skip and skim models , which we outline below . Jump based models . The method of Yu et al . ( 2017 ) reads a fixed number of words , and then may decide to jump a varying number of words ahead ( bounded by a maximum allowed amount ) in the text , or to jump directly to the end . The model uses a fixed number of total allowed jumps , and the task for the network is therefore to learn how best to spend this jump budget . The decision is trained using reinforcement learning , with the REINFORCE algorithm ( Williams , 1992 ) , where the reward is defined based only on if the model predicts correctly or not . Thus , the reward does not reflect how much the model has read . The model of Fu & Ma ( 2018 ) also has a fixed number of total jumps and is very similar to the work by Yu et al . ( 2017 ) , however it allows the model to jump both back and forth a varying number of words in order to allow for re-reading important parts . Yu et al . ( 2018 ) use a CNN-RNN network where a block of words is first read by a CNN and then read as a single input by the RNN . After each block is read , the network decides to either re-read the block , jump a varying number of blocks ahead , or jump to the end . The decision is trained using reinforcement learning , where both REINFORCE and actor-critic methods were tested , with the actor-critic method leading to more stable training . The reward is based on the loss for the prediction and the FLOPs used by the network to make the prediction . FLOP reduction is thereby directly tied into the reward signal . Huang et al . ( 2017 ) propose a simple early-stopping method that uses a RNN and reads on a word level , and where the network learns when to stop . This can be considered a single large jump to the end of the text . Skip and skim based models . Seo et al . ( 2018 ) present a model with two RNNs , a “ small ” RNN and a “ big ” RNN . At each time step the model chooses either the big or the small RNN to update the state , based on the input and previous state , which can be considered as text skimming when the small RNN is chosen . This network uses a Gumbel softmax to handle the non-differentiable choice , instead of the more common REINFORCE algorithm . Campos et al . ( 2018 ) train a LSTM that may choose to ignore a state update , based on the input . This can be considered as completely skipping a word , and is in contrast to skimming a word as done by Seo et al . ( 2018 ) . This network uses a straight-through estimator to handle the non-differentiable action choice . This approach is applied on image classification , but we include it in this overview for completeness . Other speed reading models . Johansen & Socher ( 2017 ) introduce a speed reading model for sentiment classification where a simple model with low computational cost first determines if an LSTM model should be used , or a bag-of-words approach is sufficient . The method of Choi et al . ( 2017 ) performs question answering by first using a CNN-based sentence classifier to find candidate sentences , thereby making a summary of the whole document relevant for the given query , and then using the summary in an RNN . More widely , gated units , such as GRU and LSTM , face problems with long input sequences ( Neil et al. , 2016 ) . Speed reading is one way of handling this problem by reducing the input sequence . In contrast , Neil et al . ( 2016 ) handle the problem by only allowing updates to part of the LSTM at a current time point , where an oscillating function controls what part of the LSTM state can currently be updated . Cheng et al . ( 2016 ) handle this by using memory networks to store an array of states , and the state at a given point in time comes from applying an attention mechanism over the stored states , handling the issues of older states being written over . 1https : //github.com/Varyn/Neural-Speed-Reading-with-Structural-Jump-LSTM Overall , the above state-of-the-art models are either jump or skip/skim based . We present the first speed reading model that both jumps and skips . Furthermore , current jumping-based models use a variable jump size for each input , without considering the inherent structure of the text . In contrast , our model defines jumps based on the punctuation structure of the text . This combined approach of both skipping and jumping according to text structure yields notable gains in efficiency ( reduced FLOPs ) without loss of effectiveness ( accuracy ) . We next present our model . 3 STRUCTURAL-JUMP-LSTM MODEL . Our Structural-Jump-LSTM model consists of an ordinary RNN with LSTM units and two agents : the skip agent and the jump agent . Each of these agents compute a corresponding action distribution , where the skip and jump actions are sampled from . The skip agent chooses to either skip a single word , thereby not updating the LSTM , or let the LSTM read the word leading to an update of the LSTM . The jump agent is responsible for jumping forward in the text based on punctuation structure ( henceforth referred to as structural jumps ) . A structural jump is a jump to the next word , or the next sub-sentence separator symbol ( , ; ) , or the next end of sentence symbol ( . ! ? ) , or to the end of the text ( which is also an instance of end of sentence ) . The purpose of using two agents is that the skip agent can ignore unimportant words with very little computational overhead , while the jump agent can jump past an unimportant part of the text . As both the skip and jump agent contribute to a reduction in FLOPs ( by avoiding LSTM state updates ) , the Structural-Jump-LSTM is faster at inference than a vanilla LSTM . Figure 1 shows an overview of our model : The input in each time step is the previous actions of the skip agent ( S ) , of the jump agent ( J ) , and of the current input . The output from the previous LSTM is used in combination with the input to make a skip decision – if the word is skipped , the last LSTM state will not be changed . From this we use a standard LSTM cell where the output is fed to the jump agent , and a jump decision is made . Both agents make their choice using a fully connected layer , with a size that is significantly smaller than the LSTM cell size , to reduce the number of FLOPs by making the overhead of the agents as small as possible . Section 3.1 details how inference is done in this model , and Section 3.2 presents how the network is trained . 3.1 INFERENCE . At a given time step t , Structural-Jump-LSTM reads input xi ∈ Rd , and the LSTM has a previous output ot−1 ∈ Rm and state st−1 ∈ Rm . At time step t − 1 the skip agent first takes action askipt−1 sampled from the skip-action distribution pskipt−1 and the jump agent takes action a jump t−1 sampled from the jump-action distribution pjumpt−1 . If a skip t−1 is to skip the word , then the jump agent takes no action , i.e . no jump is made . At time step t the network first needs to sample askipt from p skip t , which is computed in each time step as : pskipt = softmax ( dLIN ( state skip t ) ) ( 1 ) stateskipt = dReLU ( xt ot−1 onehot ( a skip t−1 ) onehot ( a jump t−1 ) ) ( 2 ) where dactivation is a fully connected layer with the given activation , ReLU is the Rectified Linear Unit , LIN is the linear activation , and denotes concatenation of vectors . At inference time the action can either be sampled from the distribution , or chosen greedily , by always choosing the most probable action . If the action askipt = 0 , we skip the word and set ot = ot−1 and st = st−1 , and the network will move to the next word at position i + 1 . If askipt = 1 , the word is read via the LSTM . The output and new state of the RNN is calculated and produces ot and st for the next step . The probability distribution pjumpt from which action a jump t is sampled from is computed as : pjumpt = softmax ( dLIN ( state jump t ) ) ( 3 ) statejumpt = dReLU ( ot ) ( 4 ) If the sampled action ajumpt corresponds to e.g . a jump to the next sentence , then the current LSTM output and state will be kept , and all following inputs will be ignored until a new sentence begins . When there are no more inputs the output of the RNN is used to make a final prediction . If the action is to jump to the end of the text , then the final prediction will be made immediately based on the current output .

Summary: The paper proposes a Structural-Jump-LSTM model to speed up machine reading, which is an extension of the previous speed reading models, such as LSTM-Jump, Skim-LSTM and LSTM-Shuffle. The major difference, as claimed by the authors, is that the proposed model has two agents instead of one. One agent decides whether the next input should be fed into the LSTM (skip) and the other determines whether the model should jump to the next punctuation (jump). The sentence-wise jumping makes the jumping more structural than models like LSTM-Jump, while the word-wise skipping operation has a finer skimming decision. The reinforcement learning algorithm in this paper is also different from LSTM-Jump, where LSTM-Jump uses REINFORCE, while this paper applies actor-critic approach.