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 presents a novel model for neural speed reading. In this new model, the authors combined several existing ideas in a nice way, namely, the new reader has the ability to skip a word or to jump a sequence of words at once. The reward of the reader is mixed of the final prediction correctness and the amount of text been skipped. The problem is formulated as a reinforcement learning problem. The results compared with the existing techniques on several benchmark datasets show consistently good improvements.