Explain My Surprise: Learning Efficient Long-Term Memory by predicting uncertain outcomes
Authors: Artyom Sorokin
Nazar Buzun
Leonid Pugachev
Mikhail Burtsev
Venue: NeurIPS 2022 Accept (2022)
Keywords: Memory
RNN
Information Theory
Reinforcement Learning
POMDP

Abstract: In many sequential tasks, a model needs to remember relevant events from the distant past to make correct predictions. Unfortunately, a straightforward application of gradient based training requires intermediate computations to be stored for every element of a sequence. This requires to store prohibitively large intermediate data if a sequence consists of thousands or even millions elements, and as a result, makes learning of very long-term dependencies infeasible. However, the majority of sequence elements can usually be predicted by taking into account only temporally local information. On the other hand, predictions affected by long-term dependencies are sparse and characterized by high uncertainty given only local information. We propose \texttt{MemUP}, a new training method that allows to learn long-term dependencies without backpropagating gradients through the whole sequence at a time. This method can potentially be  applied to any recurrent architecture.  LSTM network trained with \texttt{MemUP} performs better or comparable to baselines while requiring to store less intermediate data.

Explain My Surprise: Learning Efficient Long-Term Memory by Predicting Uncertain Outcomes
Artyom Sorokin
Nazar Buzun AIRI Moscow, Russia
Leonid Pugachev MIPT Dolgoprudny. Russia puleon @ maiil.u
Mikhail Burtsev
AIRI, MIPT
AIRI, MIPT Moscow, Russia burtsev burtsev@airri @airi.net
Moscow, Russia asorokin @ airi.net
buzun @airi.net
Abstract
In many sequential tasks, model needs to remember relevant events from the distant past to make correct predictions. Unfortunately, a straightforward appli- cation of gradient based training requires intermediate computations to be stored for every element
sequence. This requires to store prohibitively large inter- mediate data if sequence consists of thousands or millions elements, and as a result, makes learning of very long-term dependencies infeasible. However, usually be predicted by taking into account
the majority of sequence elements only temporally local information. On the other hand, predictions affected by long-term dependencies are sparse and characterized by high uncertainty given only local information. We propose MemUP, a new training method that allows learn long-term dependencies without backpropagating gradients through the whole sequence at a time. This method can potentially be applied to any recurrent architecture. LSTM network trained with Me performs better or comparable to baselines while requiring to store less intermediate data.
Introduction
Two dominating approaches for memory augmentation in deep learning include recurrent neural networks 2 4 7 & 9 and Transformers [T
15). Historically, recurrent networks and Transformers have been primarily developed for applications in natural language processing, which is characterized by strong dependencies between closely located elements. On the other hand, in the reinforcement learning setting, the agent often has to remember only few bits of information but for very long time [16].
Evolution of neural architectures tackling the problem of learning long-term dependencies follows the path of increasing size and complexity [77 For example, Transformers [8] and memory augmented neural networks (7] perform much better than the classical LSTM in tasks with long-term temporal dependencies. Still, even gigantic language models like GPT-3 [19] process sequences only up to 2048 elements due to quadratic complexity of self-attention. So, these architectures are much more demanding in terms of the memory capacity and computational power required for training.
To learn a long-term dependency between a distant observation the past and agent's 's actions in the current state, both the recurrent networks trained with back propagation through time and Transformers need to store gradients for all intermediate steps. Taking into account, that the minimum distance between useful information in the past and the moment of its utilization in thousands or millions of time steps, such solutions seem unrealistic given the current hardware capabilities. We propose training method that allows agent to find and store dependencies between temporarily distant events without the need to process intermediate steps.
be measured
Our contributions are the following:
36th Conference on Neural Information Processing Systems (NcurIPS 2022).

---

We propose new formulation for the problem of memory learning maximisation of conditional mutual information between memory state and future outcomes.
We introduce an original method MemUP (Memory for Uncertainty Prediction) for efficient long-term memory training by minimizing uncertainty of predictions for states where local information is not enough.
We show that implementation of our method for recurrent neural networks able to learn long-tern bongttrmm dependencies better than baselines while propagating gradients only through a small fractions of a sequence in one optimization step. Notably, the applications of UP are successful in supervised learning tasks with different loss functions well as in reinforcement learning tasks.
2 Memory to explain surprise
The main idea of our approach is very simple. Find the most surprising events or elements of a sequence and train a memory to detect and store information that allows explain out these events. Below we formulate this idea in the framework of information theory.
Lets consider sequence 7 with inputs {= and targets {yr T A model is trained to predict targets given inputs. An ideal memory could just store all previous inputs {=; ) at each step and predict Y/ from them, but is too computationally expensive. In practice memory model go has to selectively store some aggregation of this sequence a memory state m, ga (T,, m,-1) Then should be decided what to write and what to remove from the memory.
Estimating past information importance. Consider an arbitrary time step The past element ay has useful information about the target yı. if yı depends on 2t given current Tk. i.e. p(yu//a xx) o p(yk(zz) In this case, we can say that there is a temporal dependency with the length of k steps, that starts at step and ends at step k. The strength of temporal dependency remembering ze can measured as an amount of change in the distribution of yle given a knowledge of z, Multiplication of the both parts of the incquality above by p(z, (xx) and application of KL-divergence as a measure of the difference between distributions leads to the conditional mutual information measure of xy importance:
importance of
[yy; xz| = Ex KCCp(ye z (Tx) (Pp() (xx))((r((=)))))
(I)
where 24 is taken from sequences stored in some dataset generated by agent over interaction with the environment.
Objective function for memory training. Then, the problem of learning what to store in the memory can be defined as maximization of the mutual information between memory states and prediction targets with respect to the parameters #:
T
max Σ
T(y:: m. = (0 m)-1)(ax).
(2)
The sum in eq. 2 allows us to process each memory update in separate gradient step. As accounting for all future steps prevents loss of information that can be helpful in the distant future. Unfortunately. this still requires the entire remaining (future) sequence be processed for every update and has the complexity O(T') with respect to the sequence length.
Reducing the cost of memory training. Fortunately, updating the memory for all future steps is not necessary for the majority of real-world tasks. As the distance between events increases, the number of dependencies also decreases. This phenomenon is sometimes called locality of reference 20). In other words, information from the greater part of past inputs may be useless for yk prediction given dg. Accordingly, most of the elements in the sum (eq. 2 will be close to zero. Therefore, depending on the structure of the problem, we could vary the number of elements the sum, leaving only those for which long-term memory is critical for correct predictions. The main problem is to find such steps in the sequence.
Lets assume that the ideal memory m, that stores all possible information from the past is available. Then the conditional mutual information at a step t:
(y,, ".mmmmaaa m H (v (z,) H (yr xt),
(3)

---

shows the potential utility of memory for prediction at this step. When the number of elements in the sum in eq. iis is minimized to reduce computations, the steps with the highest value of mutual information from eq. 3 should be preserved to assure good quality of prediction. Naturally, we cannot estimate T(y,, m (x)) directly, because my is not available to us.
As seen from the eq 3 m, participates only in the second entropy, which indicates how well one can predict the future in the presence of ideal memory and ideal model. On the other hand, the local en- tropy H (x) can be easily estimated. A value of local entropy would be sufficient for our task if we could use it to determine the ordering for values of the potential memory utility [ye, mi (24). Specif- ically, for every pair of elements i, j if H (3)) This becomes true under one of two possible conditions: (1) H (Be(ar) > H [ z²) or (2) (ye im .z) ~ c. The first condition means that given perfect memory and a model, one could make high-quality predictions for the task at hand. The second condition means that the real distribu- tion of the target variable y/ has a similar amount of noise at each step.
H (Vzررر) then [ m. (x4) > [Vy, my (my).
If our task satisfies either one of these two conditions, then valuable information from the past could make the biggest contribution for the elements with the highest values of average "surprise" Hy (zte.
We should note that estimation of [y, mb]æz) by the local entropy makes MemUP vulnerable to the Noisy-TV problem (21), when the conditions we have described are not satisfied. Sensitivity tests for the Noisy-TV problem are discussed in Section 7
Thus, instead of optimizing the memory for each future step, we introduce new objective function that allows us to train only for steps where the potential gain from the past information is maximal:
max Σ
[ye; my = (0
KELL
here U4 denotes the set of top-K K indices of steps from t to T with the highest estimated local entropy H Hy (()])))) The hyperparameter K controls fraction of sequence to be processed in one gradient step. On the other hand, maximizing the mutual information between the memory my and an arbitrary distant event step k allows to learn long-term dependencies.
Optimization. Directly optimizing mutual information can challenging task. In our case, we use variational lower bound on mutual information proposed by [22]. Using the lower bound we can maximize mutual information by minimizing Cross-Entropt (CE) between the empirical and model distribution for all selected high entropy events:
min H.
Σ kEM
|- log (yk mm].
(S)
where mf is shortcut for m go (2e, m1-1). ds is predictor network with parameters d, that estimates probability of y given Iy and m, For a detailed derivation of eq. | see Appendix A in Supplementary Materials.
Collection of elements with high uncertainty U4 requires special uncertainty detector de model. The main requirement for de is to produce uncertainty estimates 8, = (xz)... In the simplest case, the detector estimates "surprise" log p(ve (z+)) (prediction error), which can be seen a single point estimate of uncertainty. Other models that directly estimate the uncertainty or a whole distribution can also be used.
23 24
3 MemUP for Recurrent Neural Architectures
Recurrent neural net implementation of MemUP consists of (1) training of uncertainty detector model dy, (2) selection of elements with highest uncertainty 24 and (3) training memory go and predictor 4s models. A pseudocode for RNN MemUP training is shown in Appendix G.
Uncertainty detector training. Information maximization reasoning does not limit the choice of solutions for uncertainty detector dz. In our experiments on algorithmic tasks we use a recurrent neural classifier trained with Cross-Entropy Loss. In this case detector's "surprise"" - log de (yr (x.)
is used an uncertainty estimate.
3

---

Targass
30
21
y>
3--1
yY
High uncerisinty event
Importent information observed
backproocgation: off
M
backgropagation: on
4
CE: Loss
mA0
Ro
Inputs
Ao
F1
TEPTT Length
Evenit selecte with unceriaiiiy delector dy
Figure 1: MemUP gradient update. The recurrent memory ge processes sequence. The number of prediction targets = and TRPIT length r = 2. Blue circles denote states with the highest uncertainty estimates (on ). Here æ2 marked by red circle contains information that can help to predict outcomes yr-2 and yr. At the end of RNN TEPTT rollout these states are selected to form set U3 {T - 2, T) Then, ma. 2r-2, -r-2, IT, yr are used to compute CE-Loss according eq. [ While states 2 and could separated by many millions steps, MemUP allows memory capture utility of information from 22 by propagating gradients only through the sequence of K + r elements at a time.
In Reinforcement Learning experiments we use distributional algorithm QRDQN 24 trained to predict discounted future returns R. (yt = R. in RL experiments). As QRDQN approximates whole distribution of the future returns, estimating uncertainty becomes a simple task.
Processing sequence with the memory network. MemUP allows learn long term-dependencies without propagating gradients through the whole sequence. To demonstrate that in our experiments we use Truncated BPTT train memory network ge. In the majority supervised learning experiments we set rollout length between 10 and 60 steps. In all Reinforcement Learning experiments r = 1, i.e. recurrent memory is trained without actually using backpropagation through time. All Experiments Go is a simple LSTM network with additional input encoder.
Selecting states with high uncertainty. Given uncertainty estimates {si]T=r for each future element in the sequence, we use softmax sampling p(k) e**/T/ (e*//= without replacement with 0.02. To implement sampling without replacement Gumbel-Max Trick (25) is used. A scheme illustrating a single gradient update is shown in Figure
4 Related work
The problem learning long-term dependencies has long been studied both supervised learning 226
27 and reinforcement learning domains (28). However, most of the research related to recurrent
networks architectures is aimed at solving the problems exploding and vanishing gradients 22 29)
7 30)
The common trend is that many of the existing solutions to deal with vanishing gradients
simultancously increase the cost associated with the calculation of these gradients [29 77. For
example, Transformers[10 generally require store O( of intermediate results for training sequences of length N. The recently proposed Linformer 311 and BigBird (20] architectures allow training Transformers with linear complexity in space. However, most of these new Transformers have not yet found their way in the reinforcement learning setting.
In the field of deep reinforcement learning, the dominant approaches to implementation of memory are based on recurrent neural networks combination with modern reinforcement learning algorithms (5)
32
There are many studies that propose memory architectures for specific features of the reinforcement learning problems 9 8). In the paper by Parisotto and Salakhutdinov [35). a
33
34
specifically modified DNC [7] architecture learned long-term dependencies of several hundred steps in the setting of 2D and 3D navigation tasks. The downside this architecture that the agent needs access to the information about its absolute relative location in a 2D/3D environment. In another work, Ha and Schmidhuber (36] used a complex procedure with pre- pre-training of memory and state embeddings. Their overall learning procedure is similar ours. However, their memory/world model is trained to make predictions about the next step observations only, while we train memory to make predictions about events that can be arbitrarily far in the future.
Wayne at 17 demonstrated the state the art results with the MERLIN algorithm. MERLIN's memory kept observation embeddings trained with complex variational autoencoder in DNC-

---

like 77 list. The memory module was also trained separately from the policy. Hung et al. [37 followed up on the MERLIN architecture's 's success and used the soft attention mechanism over all past embeddings to encourage exploratory actions in the POMDP environment. Their experiments showed models learning temporal dependencies with length of 500+ steps. In contrast to our work, both of these solutions still require processing an arbitrarily long sequence of state embeddings by the policy networks at each step.
The direct attempts in applying Transformer architecture in reinforcement learning setting resulted in conclusion that Transformers are too unstable to work properly in RL [38]. Parisotto et [[8] were able to overcome the instability of Transformers, showing state of the art result in the memory dependent 3D tasks. Another recent work, by Loynd et al. [39), has also applied Transformers to the POMDP problems, but the WMG architecture relies on availability of factored observations to the agent which imposes additional requirements on the environment. Switching focus from Transformers and complex architectures, Beck et al. [I6] proposed AMRL architecture: a simple modification the LSTM architecture combat the problem of learning long-term dependencies in
environments with significant amount of observational noise.
Supervised Learning Experiments
For evaluation and comparison of our method we use four tasks: Copy [40, Scattered copy, Add and permuted sequential MNIST (pMNIST) [1] All these tasks are benchmarks that are used for testing models with long-term memory. In the original Copy task sequence of length (1 = 10 out experiments) should be copied after a T - I steps (T E {120, 520, 1020, 5020) in our experiments), and in Scattered copy task a model has to make predictions in locations inside range [t, T' 1] that are chosen at random. We add this task in order to make detection of high-uncertainty locations harder. For Copy, Scattered copy and Add tasks train and test datasets have sizes 10K and 1K, for PMNIST 60K and 10K. For detailed tasks description see Appendix B.
In these experiments predictor 9a consists of three-layered MLP and a recurrent input encoder. The encoder E consists of single fully-connected layer followed by two LSTM-layers with 128 hidden units and dropout probability 0.1, Eran used to encode bigger local context than original inputs corresponding to indices in the U, set (it is also possible to use feedforward network for this task). Memory ge has the same architecture as rrss with separate weights.
We do not train a separate detector dy. Uncertainty at each step is estimated by prediction error from the predictor module qa. For more information you can look at our implementation llink in Appendix E).
In the Add Task the model is trained with MSE loss, which is special the model- -astll-assumption that target distribution a Gaussian with unit variance. Both memory and encoder trained with the same Truncated BPTT length. Truncation length equals 10 in all tasks with sequence length < 1000, 20 in tasks with length 2 1000 and 30, 60 in PMNIST 784, 3136 respectively.
of CE loss under
We compare MemUP with LSTM [2] and Transformer [10 models from pytorch library, as well as recurrent network SRNN [30] which is designed for training on long sequences and has reduced saturation of gradients. LSTM* and SRNN* have the same architecture as gran with 3 layered
MemUP
LSTM
SRNN
SRNN
Transformer
Transformertt
Copy
Scattered copy
DMNIST
Add
100
10a
10a
10-1
MSE
10
20
10
10-1
10
10-4
10-4
100
500
7000
5000
102
500
100D
5000
7B4
3136
150
500 1000
5000
sequences length/T)
Soquences longthiT)
sequances length(T)
sequences length/))
Figure 2: Final performance on supervised learning tasks. X-Axis shows task scaling, while Y-axis correspond to metric score at the end of training. Metrics: Inverted Accuracy (I. Accuracy) in Copy, Scattered copy and PMNIST tasks, MSE in Add task. All curves are averaged over runs.
5

---

MLP on top. In the case of SRNN baselines, we have replaced all LSTM layers with SRNN layers. The Transformer model consists of a single FC layer followed by 4 self-attention layers and the same 3 layered MIP on top. Baselines were trained in two modes. In the first mode (columns with Transformer, SRNN, LSTM) all elements of input sequence of length T were fed into a model simultaneously, so gradients propagate through the whole sequence during the training. We don't t have results for Transformer on sequences of 5000+ steps since it didn't fit into our GPU RAM. In the second mode (Transformer Tr., SRNN Tr., LSTM Tr.) BPTT rollout for recurrent architectures and attention window for Transformer were truncated to process the same number of steps as MemUP (r + K).
Results in Figure 2 show that only MemUP allows to achieve good quality of prediction in all tasks. In truncated BPTT setting (solid lines) Mem is significantly better than LSTM tr, SRNN tr or Transformer tr due to MemUU ability to learn temporal dependencies longer that the length of BPTT rollout. When full information about the sequence is available to LSTM, SRNN and Transformer (dashed lines) but not MemUP, the latter still shows the best results for longest versions of Copy, Scattered Copy, pMnist, and has better overall performance. Moreooeer, when compared with Full BPTT baselines, MemUP requires to store much less intermediate computations (10-250 times) due to usage of short Truncated BPTT rollouts.
Reinforcement Learning Experiments
Performance of the MemUP algorithm is studied in two reinforcement learning tasks. The first task is classic T-Maze environment where the main difficulty is a long-term dependency between hint at the starting position and location of a reward at the exit (see Figure 3 1). The main advantage of this problem that it allows test the agent's memory mechanisms isolation because of primitive policy and the observation space. We base our experiments on the noisy version introduced by [16).
The second task is the color dependent object collection in the Doom environment 442 43 which
requires long-term memory in combination with a more complex policy and rich observation space. In the task introduced by [44 the agent is placed in room filled with acid (see Figure 2.) and constantly loses health. Objects of green and red colors are scattered throughout the environment. Object of one color replenish the agent's health and give and give a -1 reward. The correspondence between effects and colors is determined randomly at the beginning of each episode. This information is conveyed by a column whose color matches the color of health replenishing items. In our version the column disappear after first 45 steps (see Figure 3). The episode ends when the agent's health drops to zero after 1050 time steps.
reward, while others take away health
The most straightforward extension of the MemUP algorithm to reinforcement learning problem is use memory state m, an additional input to RL agent. There several possible solutions to combine MemUp's training with agent's 's policy training: (1) pre train MemUP and then train an agent, (2) use alternating memory and agent training phases, train memory and agent in parallel, (3) combining the MemUP and RL agent into single into an end-to-end network. In these experiments, we test the simplest version: train McmUP on trajectories generated with a frozen policy, then train an agent with a frozen MemUP memory. The second version is also implemented in the code (link in Appendix E).
Hint
Final Choice
R
Agent
Figure 3: Environments for testing long-tern time dependencies. 1. T-Maze environment. The agent should reach the T-shaped junction and choose one of the arms (L or R). A hint about what arm choose is provided at the very beginning. 2. and 3. Vizdoom-Twoo Two-Colors environment. The agent is in the room and constantly loses health. To replenish his hcalth and recieve a reward the agent needs to collect items of the same color as the column. The color of the column is chosen randomly from red green options at the beginning of the episode. After 45 steps from the beginning the column disappears. The episode lasts for maximum of 1050 steps.

---

Memtil (ec()))
PPOLSTM--100
PPOLSTM)-S03
AMER-1N-020 AMELL(1-1000)
ANILL 123
prevekingf-1)
AMALL01100)
AMRL-WAX00)
Mean Episode Rewaid
900
2
Mean Eptcods
20
20
Number of Steps
Number of steps
(a)
(b)
a) Learning curves for MemUP and baselines in T- T-MMaz-LNR--OOE LNR-1000. b) Learning plots in the Vizdoom-Ti Two-Colors environment for agents. curves are averaged over runs. Translucent areas around each curve show the standard deviation computed over all runs. The memory pre- pre-training phase of MemUP agent is marked with the dashed line. TBPTT lengthlattention span for each baseline is specified in the parenthesis.
Figure 4:
Inputs 2t are formed from observations or and previous action d--1- encoded using convolutional networks or fully connected networks. We use discounted future returns R₂ Σ "r, as targets yt- the experiments, MemUU compared with the following baselines:
This time Observations are
PPO-LSTM is recurrent version of PPO 45 with single LSTM-layer 2 We use PPO-LSTM
implementation from the RLPyt library
46
IMPALA-ST is an IMPALA 44 agent using Stabilized Transformer 10 architecture. Stabilized Transformer was presented by Parisotto et al. [8]. For experiments, we use the only available open
implementation of this algorithm [48]
AMRL is proposed by [6 for the Noisy T-Maze Task. In the original study AMRL was compared with many different memory architectures including DNC, multi- LSTM and outperformed all of them on T-Maze-LN-10 task as well as other memory experiments. AMRL similar to PPO-LSTM baseline, but extends LSTM with AMRL Layer.
In all experiments we the same encoder networks for observation embeddings for all baselines and MemUP. Additionally the size of an LSTM-layer is the same for the memory module PPO-LSTM and PPO-AMRL baselines. Results for the T-Maze-LNR- MiizeeLLR--000 (10 times longer dependency than in (16)) environment are presented in Figure 4a The designation T-Maze-LNR- -000
1000 means that the maze length (and corresponding tempral dependency) in each episode can be from 1000 1009 steps. Several rollout lengths are tested for CONNbased baselines. We could not successfully train single PPO-LSTMM or AMRL agent even with TBPTT rollout of 100d"]
During MemUP pre-training phase, episodes are generated with a random policy. We use the discounted future reward with 7 0 as prediction target. policy training for the MemUP agent starts after the memory pre-training phase is completed. The time spent in the pre-training phase marked with dashed line in all figures. As shown in Figure4a the agent learns almost instantly, given memory pre- pre-trained with our method. For the MemUP agent we train all components including memory module and uncertainty detector for each individual run from scratch. Additionally, Appendix Dshows that MemUp can be trained to solve Simaze-LNR-20000
Results for the Vizdoom-Two-CA Two-Colors environment shown in Figure 4b IMPALA-ST(atune with a short attention span could not learn any reasonable policy. On the other hand, IMPALA- ST(att=100) solves the problem and reliably survive in the environment for 1050 steps. Both PPO-LSTM and AMRL baselines with TBPTT rollouts of 128 and 8 steps could not learn to survive longer than a random agent. While the rollout of 128 steps potentially enough to remember the color of the column or previously collected items, PPO-LSTM and AMRL agents are unable to utilize long-term information from the past.
We tested our AMRL implementation T-Maze-LNE LNN- 100 with full rollout and get results aligning with the original paper. But the algorithm did not withstand 10 times increase in episode length and truncation constraints
7

---

For MemUP we use the discounted future reward with = 0.8 as prediction target. The memory module do is trained with the rollout length of step. Vizdoom-Two-Ge
we set number of targets K = 3. Thus, to train the memory we use 4 separate observations per episode for one gradient update. In ViZDoom, we train PPO (MemUP++PPO and QRDQN (MemUP+QRDQQN) agents with MemUP pretraining.
Both MemUP+PPO and MemUP+QRDQN learn to survive in the environement for full 1050 steps by collecting healing items. Though the final performance of IMPALA-ST(attn = 100) is slightly better. We speculate that this is due to the fact that the pre-trained memory M has learned to store mostly information about the color of the column, while IMPALA-ST (attn= 100) has slightly better spatial awareness. On the other hand, MemUP use significantly fewer resources in terms of the size of processed sequences during training/pre-trainn outperforms baselines with comparable resource requirements, like IMPALA-ST (attn=8), PPO- (S) and AMRL-MAX(=8))
and evaluation. As well, MemUP significantly
7 Sensitivity to the Noisy-TV problem
In section 2 we discussed conditions on true uncertainty of targets y/ given a perfect memory mç. The goal to set conditions under which we make good guess about relative values of potential memory utility 1 [ye c)) by estimating local entropy H (y/ c.). However, this may not be the case. For example, some targets y/ may be completely unpredictable and no information from the past can reduce the entropy for the distribution of yą: H (yllz,) (z,) ~ H (yt ( Then the local entropy is large, while potential utility of memory for Y/ predictions is close
Training memory to predict such targets is pointless and prevents MemUP from focusing on events for which training long-term memory essential. The problem of over-emphasis Noisy-TV problem 21]. Noisy-TV problem is often encountered by algorithms using curiocity-based exploration in reinforcement learning 49] To test robustness of MemUp the Noisy-TV problem. we have modified T-maze-LNR-10l environment.
unpredictable surprising events is called
Noisy T-maze With Distractors. In this version, agent can receive +4 or rewards in D+1 decision points. Their location in the corridor is chosen randomly at a distance of at least 50 steps from the hint. All of them can be detected when o, [1] = element observation vector ot. Decision points are distinguishable from each other. Each decision point has unique value of or [1]. Only in one
decision point, the next reward depends on the agent's action and the value of the hint [o,[4] 1.
in this case ). In other D decision points, the next rewards are completely random ( chosen with probability of 0.5 ). Thus, Noisy T-maze With Distractors generate sequences with one long-term dependency and D events acting as Noisy-TV distractors.
MemUP: Aupdates before solving the
1520
K-2
1220
2140
3500
1500
$2200
1200
2280
2300
7120
33400
K-5 1200
5590
C:10
1530
2000
2150
1340
11400
Figure 5: Each cell shows the mean number of updates required for the memory to correctly predict (achieving 0.1 RMSE) the reward after the decision point connected with the temporal dependency. The x-axis shows the number of unpredictable decision points in the environment. The Y-axis shows number of events predicted by MemUP. Darker colors means slower solution.
8

---

We conducted 25 experiments with 6 independent runs in each (150 runs in total). The differences in the experiments come down to two parameters: number of distracting events D (X-axis in Fig. [ and number of events sampled for prediction K (Y-axis in Fig. 5 More distracting events make the task of learning the temporal dependency harder. Increasing the size of L4, set( by changing K) improves the chance of selecting future event that is a part of temporal dependency. On the other hand, with an increase in 24, the resource cost for memory training also increase. The TBPTT rollout length in these experiments is 1.
Results in Fig- 5 show that with an increase the number of distractions, memory learning speed decreases. However, increasing the set U4 allows to alleviate this problem. For the experiment (D=19, K=1), none of the runs solved the task in the first 45000 updates. In experiments (D=19, K=2) and (D=19, K=3) only 4 out of 5 runs solved the task. The results show that the agent can learn the time dependeney fairly quickly even the uncertainty detector select noisy events 80% of the time.
8 Ablation study
To study an effectiveness of our core idea, i.e. training memory by predicting long-term future events with high uncertainty, we compared MemUP with baselines that exclude core MemUP features while sharing the same neural architecture.
In these experiments, carefully consider the process of learning memory (without further training the RL agent) on the T-Maze-LNR problems. The following memory ablation baselines were studied:
MemUP: Proposed in this paper. See section 2
Same as MemUP, but steps for future outcome prediction selected randomly and uniformly among all future timesteps. This randomization verify whether predicting steps with high uncertainty helps MemUP overcome the problem of learning long-term dependencies with small BPTT rollouts.
Rnd-Pred:
Default: Same MemUP, but it is trained to predict return R₃ at each step oposed to an arbitrarily distant future events as in MemUP and Rnd-Pred. Thus, Default is the same as the LSTM baseline in the Supervised Learing Experiments (see But in this study test this bascline on all intermediate Truncated BPTT lengths.
All versions have the same architecture, and hyperparametes. They were trained in T-Maze-LNR-100 and T-Maze-LNR 1000 problems. Episodes were generated by a random strategy. All versions are trained with 5 different Truncated BPTT lengths: 1, 10, 50, 100. Testing metric is root-mean- squared error(RMSE) between the models' prediction and the actual return at the final step of the episode. The evaluation was conducted on the 100 separately recorded episodes.
Figure 6 shows the results of the experiments. On the x-axis, we have specified the length of the Truncated BPTT in decreasing order, i.c. in the order of increasing task complexity. For every rollout length there a corresponding absolute value of error between predictions of the model and targets
rndegred
Root: wwr--MMan----aa Square Emor
Root-Mean Square Error
1-50
10
1-0
[-00
1-
Truncated BPTT/Rallout Length
Truncated AoTT/o Length
(a)
(b)
Figure 6: Averaged over 6 runs prediction error the end of training for each ablation baseline with respective TBPTT length. Vertical bars show standard deviation computed over 6 runs a)
T-Maze-LNR- 1000 b) T--Maze-LNR--000

---

at the end of training (3e4 updates for T-Maze-LNR-100 100 and 3e5 for T-Maze-LNR--OOG Smaller values represent better predictions.
The Default modification is unable to learn a long-time dependency exceeding the length of the tested rollouts. This is the same the results of the PPO-LSTM agent the T-Maze-LNN Reinforcement Learining experiments.
The Rnd-Pred modification solves the task with rollouts 100 and 50 in T-Maze of length 100, but for shorter rollouts and for longer T-Maze-LNR-1000 results deteriorate significantly. In the majority of runs in T-Maze-LNR 1000, the memory has not been learned. The relative success of Rnd-Pred is expected because random selection of prediction target still allows to sample a useful future event with some small probability Thus, random sampling of prediction targets allows to learn casual dependency that is 2 times longer than TBPTT, but already struggles at ratio of 10.
On the other hand, the MemUP is able to solve the problem regardless of the rollout length, i.e. MemUP demonstrated an ability to learn casual dependencies that 1000 times longer than its TBPTT length in this setting. In this study MemUP is successfully trained with TBPTT of length 1, in other words, without using backpropagation through time at all.
Conclusion
In this paper, we proposed new method for training long-term memory. The main idea is to train memory network predict future outcomes of high uncertainty and skip all others. Predicting a small number of arbitrarily distant future outcomes substantially saves computational resources required backpropagate gradients. At the same time, the emphasis on high uncertainty outcomes allows not to miss long-term dependencies in the task.
Experimental results show that our training algorithm allows to learn temporal dependencies signifi- cantly longer than the number of steps processed for a single gradient update. None of the baseline architectures trained in classical way demonstrate such ability. All these baselines with comparable results require to store at least 200 times more intermediate calculations for the Add, Copy, Scattered Copy, PMNIST tasks as well as at least 500 times more for the T-maze, and 625 times more for the Vizdoom-Two-Coll Two-Colors (Stabilized Transformer baseline has quadratic complexity while processing 25 times more steps).
Even using fewer resources, MemUP outperforms non- con-trunccted baselines that has simultaneous access all elements of sequence (including Transformer) PMNIST(3136) Copy and Scattered Copy tasks. On T-maze and Vizdoom Soaloom-Two-CGe Two-Colors tasks, MemUP is better than all recurrent baselines. Another advantage is that MemUP can be combined with any recurrent architecture and applied both to supervised as well reinforcement learning settings.
We believe that the MemUP algorithm applied to recurrent networks has shown promising results, and the main idea demonstrates exciting avenues for future research.
10 Acknowledgments and Disclosure of Funding
This work was partially supported by grant for research centers in the field of artificial intelligence, provided by the Analytical Center for the Government of the Russian Federation accordance with and the agreement with the Moscow Institute of Physies and Technology dated November 1, 2021 No. 70-2021-00138. This work was partially supported by the Ministry of Science and Higher Education of the Russian Federation (Grant No. 075-15-2020-8011
the subsidy agreement (agreement identifier 000000073032
References
[1] Paul Werbos. Backpropagation through time: what it does and how to do it. Proceedings of the IEEE, 78(10):550-1566) 1560, 1990.
[2] Sepp Hochreiter and Jürgen Schmidhuber. Long short-term memory. Neural computation, 9()): 7785-1780, 1997.
For example, for rollout of 50 steps, the chance select the correct target step is 0.5
10

---

[3] Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. arXiv:1410.5401 2014.
[4] Jason Weston, Sumit Chopra, and Antoine Bordes. Memory networks. arXiv:1410.39/6 2014.
[5] Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lilli- crap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In International conference on machine learning, pages 1928-1937, 2016.
[6] Junhyuk Oh, Valliappa Chockalingam, Satinder P. Singh, and Honglak Lee. Control of memory, active perception, and action in minecraft. In Faoria- Balcan and Kilian Q. Weinberger, editors, Proceedings of the 33nd International Conference on Machine Learning, ICML 2016, New York City, NY, USA, June 19-24, 2016, volume 48 of JMLR Workshop and Conference Proceedings, pages 2790-2799, J.rR. 2016.
.ml...ggmm........mmmmmm..m.mm.. mlr press/
1y48/oh1..np html
[7] Alex Graves, Greg Wayne, Malcolm Reynolds, Tim Harley, Ivo Danihelka, Agnieszka Grabska- Barwiñisk,, Sergio Gómez Colmenarejo, Edward Grefenstette, Tiago Ramalho, John Agapiou, et al. Hybrid computing using neural network with dynamic external memory. Nature, 538 (7626):471, 2016.
[8] Adam Santoro, Ryan Faulkner, David Raposo, Jack W. Rae, Mike Chrzanowski, Theo- phane Weber, Daan Wierstra, Oriol Vinyals, Razvan Pascanu, and Timothy P. Lilli- crap. Relational recurrent neural networks. In Samy Bengio, Hanna M. Wallach, Hugo and Roman Garnett, editors, Advances Neural Informa- tion Processing Systems 2018, NeurIPS 2018, December 3-8, 2018, Montréal, Canada, pages 7310-7321, 2018. URL https: ppwoa//d/nn/ceeeegs//eee newripe hccpppppe/211//aa//// 22cebbtf637722000eeedddffo37/723 2tt.://5s56-Abbtaaat.... Abstract.hn html
Larochelle, Kristen Grauman, Nicolò
in Neural Information Processing Systems 31: Annual Conference
[9] Steven Kapturowski, Georg Ostrovski, John Quan, Remi Munos, and Will Dabney. Recurrent experience replay in distributed reinforcement learning. 7th International Conference on Learning Representations (ICLR), 2019.
[10] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Isabelle Guyon, UI- rike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett, editors, Advances in Neural Information Processing Systems 30: An- nual Conference on Neural Information Processing Systems 2017, 4-9 December 2017, Long Beach, CA, USA, pages 5998-6008, 2017. URL URL http::/paperrs. petpapers.niss..... nips cttppper///p/////////pppppp 7181-attention is- allaalllyy neee
[11] Mikhail S Burtsev, Yuri Kuratov, Anton Peganov, and Grigory Sapunov. Memory transformer. arXiv preprini arXiv:2006. 2020.
[12] Jie Lei, Liwei Wang, Yelong Shen, Dong Yu, Tamara Berg, and Mohit Bansal. MART: Memory- augmented recurrent transformer for coherent video paragraph captioning. In Proceedings of the 58th Annual Meeting of the Association for Computational Linguistics, pages 2603-2614 14, Online, July 2020. Association for Computational Linguistics. doi: 10. 10..18653/11122220acc.mmaaa 233. URL httpittts///aaaaattollooo
www.or/gg........./ acm-main..m..mmmmmmmmmmmmmm.mm.m 233
[13] Ankit Gupta and Jonathan Berant. Gmat: Global memory augmentation for transformers. arXiv preprint 2000:20006.3227,, 2020.
[14] Qingyang Wu, Zhenzhong Lan, Jing Gu, and Zhou Yu. Memformer: The memory-augmented transformer. arXiv preprint arXiv.2010.06891, 2020.
[15] Angela Fan, Thibaut Lavril, Edouard Grave, Armand Joulin, and Sainbayar Sukhbaatar, Address- ing some limitations of transformers with feedback memory. arXiv preprint Xiv:2002.09402. 2020.
11

---

[16] Jacob Beck, Kamil Ciosek, Sam Devlin, Sebastian Tschiatschek, Cheng Zhang, and Katja Hofmann. AMRL: aggregated memory for reinforcement learning. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URI. httptttt///ppaaeevieee net/EooruniddBkk7750a
[17] Greg Wayne, Chia-Chun Hung, David Amos, Mehdi Mirza, Arun Ahuja, Agnieszka Grabska- Barwinska, Jack W. Rae, Piotr Mirowski, Joel Z. Leibo. Adam Santoro, Mevlana Gemici, Malcolm Reynolds, Tim Harley, Josh Abramson, Shakir Mohamed, Danilo Jimenez Rezende, David Saxton, Adam Cain, Chloe Hillier, David Silver, Koray Kavukcuoglu, Matthew Botvinick, Demis Hassabis, and Timothy P. Lillicrap. Unsupervised predictive memory in a goal-directed agent. CoRR, abs/ 888.100660 2018. URL http::/axxx:/i..../ org/ abs/ 1803.
[18] Emilio Parisotto, H. Francis Song. Jack W. Rae, Razvan Pascanu, Çaglar Gülçehre, Siddhant M. Jayakumar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark, Seb Noury, Matthew M. Botvinick, Nicolas Heess, and Raia Hadsell. Stabilizing transformers for reinforcement learning. CoRR, abs/1910.06764, 2019. http:////aax
[19] Tom B Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. arXiv preprint arXiv:2005. 140:: 2020.
[20] Manzil Zaheer, Guru Guruganesh, Kumar Avinava Dubey, Joshua Ainslie, Chris Alberti, Santiago Ontañón, Philip Pham, Anirudh Ravula, Qifan Wang, Li Yang, and Amr Ahmed. Big Transformers for longer sequences. In Hugo Larochelle, Marc' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Tien Lin, editors, Advances in Neural Information Processing Systems 33: Anmual Conference Neural Information Processing Systems 2020, NeurlPS 2020, December 62, 2020, virtual, 2020. URL httph::t///pppe p/tope//// wwutinngg..euuppppe neur ips cc/ ppaaper/22000 hash/ (25500014222844775533333333 a-6aab9-36111----- Abstract.
bird:
[21] Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov. Exploration by random network distillation. arXiv preprint arXiv: 1810. /2894, 2018.
[22] David Barber Felix Agakov. The im algorithm: a variational approach to information maximiza- tion. Advances in neural information processing systems, 16(320)2201, 2004.
[23] Evgenii Tsymbalov, K. Fedyanin, and Maxim Panov. Dropout strikes back: Improved uncer- tainty estimation via diversity sampled implicit ensembles. ArXiv, 0bs/2003.032777 2020.
[24] Will Dabney, Mark Rowland, Mare G Bellemare, and Rémi Munos. Distributional reinforcement learning with quantile regression. In Vhirot-Second AAAI Conference on Artificial Intelligence, 2018.
[25] Iris AM Huijben, Wouter Kool, Max B Paulus, and Ruud JG van Sloun. A review of the gumbel-max trick and its extensions for discrete stochasticity in machine learning. arXiv preprint arXiv:2110.01515, 2021.
[26] Herbert Jaeger. Tutorial on training recurrent neural networks, covering BPPT, RTRL, EKF and the" echo state network" approach, volume 5. GMD--Forchungszeentuu Informatiinnsechnik Bonn, 2002.
[27] Sepp Hochreiter, Yoshua Bengio, Paolo Frasconi, Jürgen Schmidhuber, et al. Gradient flow in recurrent nets: the difficulty of learning long-term dependencies. A field guide dynamical recurrent neural networks. IEEE Press, 2001.
(28) Bram Bakker. Reinforcement learning with long short-term memory. In Thomas G. Dietterich, Suzanna Becker, and Zoubin Ghahramani. editors, Advances in Neural Information Process- ing Systems 14 [Neural Information Processing Systems: Natural and Synthetie, NIPS 2001, December 3-8, 2001, Vancouver, British Columbia, Canada], pages 1475-1482. MIT Press, 2001.
[29] Alex Graves, Greg Wayne, and Ivo Danihelka. Neural turing machines. CoRR, abs/1410.5401, 2014. URLphttp://arant org/ abss/411..4pp
12

---

[30] Michael Rotman and Lior Wolf. Shuffling recurrent neural networks. CoRR, abs/2007 2020. htttthtt///aaattt wwtp:/iiww.rggbb////////// 07324
[31] Sinong Wang, Belinda Z. Li, Madian Khabsa, Han Fang, and Hao Ma. Linformer: Self-attention with linear complexity, CoRR, abs/2006.04768, 2020. URL https://aaxiv..rrg wttpiiiv..r//bb///0/6//
04768
[32] Matthew Hausknecht and Peter Stone. Deep recurrent q-learning for partially observable mdps. CoRR, abs/l507.06527,, 7(1), 2015.
[33] Alexander Sasha Vezhnevets, Simon Osindero, Tom Schaul, Nicolas Heess, Max Jaderberg, David Silver, and Koray Kavukcuoglu. Feudal networks for hierarchical reinforcement learning. In Doina Precup and Yee Whye Teh, editors, Proceedings of the 34th International Conference on Machine Learning. ICML 2017, Sydney, NSW, Australia, 6-11 August 2017, volume 70 of Proceedings of Machine Learning Research, pages 3540-3549. PMLR, 2017. / //proceedings.ni. mir. htt.://g....iii.peeeeee///////e/ pttpp///////ee/e///eeeeeeeeeeeee
html
[34] Bowen Baker, Ingmar Kanitscheider, Todor M. Markov, Yi Wu, Glenn Powell, Bob McGrew, and Igor Mordatch. Emergent tool use from multi-agent autocurricula. In 8th International Conference on Learning Representations, ICLR 2020, Addis Abuba, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URI. httpps:///ppnrrevie net/forum? Ses.rrrn
[35] Emilio Parisotto and Ruslan Salakhutdinov. Neural map: Structured memory for deep rein- forcement learning. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30- May 2018, Conference Track Proceedings. OpenRe- view.net, 2018. /tLLitttps//openrreee
[36] David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In Samy Bengio, Hanna M. Wallach, Hugo Larochelle, Kristen Grauman, Nicolò Cesa-Bianchi, and Roman Gamett, editors, Advances in Neural Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurIPS 2018, December 3-8, neur ips
2018, Montréal, Canada, pages 2455-2467, 2018. URL httss/tt/:/// 1a2ba Abstract.
[37] Chia-Chun Hung, Timothy Lillicrap, Josh Abramson, Yan Wu, Mehdi Mirza, Federico Carnevale, Arun Ahuja, and Greg Wayne. Optimizing agent behavior over long time scales by transporting value. Nature communications, 10(1):1-12, 2019.
[38] Nikhil Mishra, Mostafa Rohaninejad. Xi Chen, and Pieter Abbeel. A simple neural attentive meta-lcarner. In 6th International Conference on Learning Representations, ICLR 2018, Van- couver, BC, Canada, April 30 May 3, 2018, Conference Track Proceedings. OpenReview.net, 2018. URL httpsttps:/// btttp::://ppnneeveee
net/
[39] Ricky Loynd, Roland Fernandez, Asli Çelikyilmaz, Adith Swaminathan, and Matthew J. Hausknecht. Working memory graphs. In Proceedings of the 37th International Con- ference on Muchine Learning, ICMI, 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 6414-6414 PMLR, 2020. URL http:///proceedings wtt.eemmrr.prrrrrr presss/119/1oyyd200 html
[40] Martin Arjovsky, Amar Shah, and Yoshua Bengio. Unitary evolution recurrent neural networks. International Conference on Machine Learning, pages 1 1200-122, 2016.
[41] Quoc V. Le, Navdeep Jaitly, and Geoffrey E. Hinton. A simple way to initialize recurrent 2015.
networks of rectified linear units. ArXiv,
[42] Michat Kempka, Marek Wydmuch, Grzegorz Rune, Jakub Toczek, and Wojeiech Jaškowski. Vizdoom: A doom-based ai research platform for visual reinforcement learning. In 2016 IEEE Conference Computational Intelligence and Games (CIG), pages 1-8. IEEE, 2016.
[43] Guillaume Lample and Devendra Singh Chaplot. Playing FPS games with deep reinforcement learning. Satinder Singh and Shaul Markovitch, editors, Proceedings of the Thirty-Firs AAAI Conference on Artificial Intelligence, February 4-9, 2017, San Francisco, California, USA, pages 2140-2146. AAAI Press, 2017. hRIptttp://// /aaii org/ocg. wwtgg//cs/ndee..../ ////ddddxx,./ andex.ppp//aee AAAI/
AAAI17 hwwp1///aapr///ee pwwpr/ htcaapr/vieeee hwtp/////////////6//
13

---

[44] Edward Beeching, Christian Wolf, Jilles Dibangoye, and Olivier Simonin. Deep reinforcement learning on a budget: 3d control and reasoning without supercomputer. CoRR, 2019. URL (RR.htp:: /tttp://arxiv 19wrg/4b/1190..... 01806
[45] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:1777003444 2017.
[46] Adam Stooke and Pieter Abbeel. rlpyt: A research code base for deep reinforcement learning in 2019.
pytorch. arXiv preprint 1909.
[47] Lasse Espeholt, Hubert Soyer, Rémi Munos, Karen Simonyan, Volodymyr Mnih, Tom Ward, Yotam Doron, Vlad Firoiu, Tim Harley, lain Dunning, Shane Legg, and Koray Kavukcuoglu. IMPALA: scalable distributed deep-rl with importance weighted actor-learner architectures. In Jennifer G. Dy and Andreas Krause, editors, Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmmsssan, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pages 1406-1415, PMLR, 2018. URL httphtttpp/prcceedinnn
mlr prw.s/essss.... wttppp/eeeeeeeeeeeeeeeeeree wsw.holthhhhtt html
[48] Shakti Kumar, Jerrod Parker, and Panteha Naderian. Adaptive transformers rl. arXiv preprint arXiv:2004.03761, 2020.
[49] Yuri Burda, Harri Edwards, Deepak Pathak, Amos Storkey, Trevor Darrell, and Alexei A Efros. Large-scale study of curiosity-driven learning. arXiv preprint
2018.
14


Reviews:
[{"review_id": "xLPeyyD4k60", "reviewer": "Reviewer_Mqk3", "summary": "The work introduces memUP, a training method that allows to learn long-term dependencies without backpropagating gradients through the entire sequence (unlike RNNs). The method is based on maximizing mutual information between memory states. Then, it is explained how the equations need and can be approximated. This introduces further assumptions leading to possible vulnerabilities to the Noisy-TV problem. This is further analyzed in the experimental sections. This is further simplified to a practical optimization function related to cross-entropy loss that depends on the memory $m$ at a timestep and the model $q$ given the respective input.\nThe work then continues by describing how memUP is applied to RNNs. The experimental sections cover results for basic sequence tasks (copy, scattered copy, permuted MNIST, and add task), while comparing to LSTMs, SRNNs, and Transformers. Also, it is experimented in RL settings (T-maze, Vizdoom-Two-colors), showing competitive results to existing methods.\n\n", "questions": "* How is the optimization of the CE loss (eq 5) is solved for both $\\theta$ and $\\phi$? When is the stop_gradient(m) applied? \n* How would the method needs to be modified to solve regression problems? (for $y_t$ scalar).\n* How harder it is to train (and converge) this memory model vs LSTMs with chrono initialization? \n* How does this method perform on more real sequential tasks such as language modeling?\n* How does the training effort changes when the tasks are more challenging and more parameters could benefit the solution? Is memUp method still viable? If not, what would need to happen?\n* What is the effect of the memory size on the solution? \n* Would it be possible to measure the effect of the memory on the quality of the solution? For example, what is the performance of the model trained with memUp but with the solution measured with $m_t$ always set to 0 (zero) to \"remove\" the memory information.\n", "limitations": "The authors describe and analyze the Noisy-TV problem based on their approximation.\nThey also assume that dependencies are scarce with the increase of the distance.\n\n", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The work seems original to be the best of this reviewer knowledge and relevant to the NeurIPS community. Enabling models to learn and process long-range dependencies is indeed a problem of interest. In particular, how to train memory augmented models. The works proposes an interesting method to train such models. \n\nThe idea behind the method is clear and interesting. The paper focuses on deriving the method, but could have make a better job exploring it further. Instead, authors decided to compare against other methods in two settings (supervised learning on sequences and RL). See questions below. In addition, several of these experiments could have compared to newer recurrent methods, and memory augmented methods (both transformers and RNNs). See [1] and [2] as examples. \n\nThe text is clear. Putting more emphasis on how the algorithm is applied to RNNs (and other networks) could really help the reader understand the method and give it more clarity. Defining the tasks in the text instead of the supplementary would be appreciated as well.\n\n\n[1] Legendre Memory Units: Continuous-Time Representation in Recurrent Neural Networks, Voelker et al., 2019\n\n[2] Not All Memories are Created Equal: Learning to Forget by Expiring, Sukhbaatar et al., 2021", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "dKScUlt46X1", "reviewer": "Reviewer_j9Pj", "summary": "The paper proposes MemUP (Memory for Uncertainty Prediction), a method to learn long-term dependencies for recurrent models. MemUP learns to keep useful past states for future use, so it saves both computation and memory. The memory module is trained to optimize the mutual information between past states and current prediction targets, and only the steps that have highest \"uncertainty\" would be used for training to further save memory usage. Empirical experiments on supervisied learning and reinforcement learning tasks show that MemUP significantly boost the results over its baselines.", "questions": "I have a few questions that are not clarified after reading the paper:\n\n1\\. In the uncertainty, how do you get p(y\\_t | x\\_t), given you are training a model to estimate the p(y\\_t | x\\_t, m_t)? Is it a separate module (L190)?\n\n2\\. For the RL experiments, you set r=1. Can you justify this, given it might hurt the performance of other models? I wonder what the results will be if r is much larger.\n\n3\\. In equation 2, the objective function for memory training assumes that the current memory m\\_t should be helpful for future prediction at step k, while the m\\_t will be updated to m_k at that moment. If my understanding is correct, there is a discrepancy between the objective and inference. Is that right?", "limitations": "I do not foresee any potential societal impacts that may occur.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The solution proposed by the paper is valuable for long-sequence modeling. On the one hand, it reduces the computational overhead of RNN model training; on the other hand, it could preserve distant information, which RNNs would easily forget. The methods are simple, and the intuition behind the methods all make sense to me. The paper is well-written, with details clearly formulated and illustrated. Together with its experiments on both supervised learning and reinforcement learning tasks, I can foresee its potential values in the field of reinforcement learning and NLP.\n\nHowever, the paper does come with some drawbacks and some confusions:\n\n1.  Lack of background work: Though closed related to memory neural networks, none of them is discussed and compared. For example, J Weston, S Chopra, and A Bordes. \"Memory networks.\" ICLR (2014). Similarly, pointer networks should also be discussed, e.g. See et al. (2017). Memory modules are also popular for long-sequence processing in NLP, e.g. coreference resolution (Toshniwal et al. 2020).\n    \n2.  Insufficient empirical evidence: Experiments in section 5 are claimed to be long-sequence tasks, while their sequences are not long enough. The lengths of Copy, Scattered copy, and Add are 1000, which can even be handled by some transformers (e.g. T5). The truncation lengths are as small as 10 and up to 60 only, which can be much larger for real applications.\n    \n3.  Choice of uncertainty detector: While the description of \"Reducing the cost of memory training\" in section 2 seems intriguing, the method used in training is simplified into a point estimate of prediction error. In this case, it makes more sense to me to simplify the justification in section 2.\n\n5.  While MemUP uses information theory to help the model to keep the most useful information, the basic idea of MemUP is very similar to LSTM, which purely relies on BPTT to learn how to keep memories. The authors show in sections 5 and 6 that MemUP outperforms LSTM, the baselines are presented in a negative light where the rollout lengths are less than 60 steps. The authors may clearly compare the difference between MemUP and LSTM, as well as providing more evidences to show the advantage of MemUP.  ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ""}, "extra_text": {}}, {"review_id": "K0f78_L86a3", "reviewer": "Reviewer_w1z8", "summary": "In this paper, the authors proposed a method called MemUP to help recurrent neural networks learn long-term dependencies better. Specifically, the authors leverage a memory model to predict future outcomes with high uncertainty. By skipping all states with lower uncertainty, the training process takes less computational cost in backpropagation. The framework involves a uncertainty detector model and a predictor to estimate mutual information used in memory model training. The authors show experimental results on both supervised learning tasks and reinforcement learning tasks with improvement comparing with baselines.", "questions": "1. What is the computational cost of the proposed method?\n2. By involving auxiliary models to estimate \"surprise level\", how accurate is the estimation in the proposed framework and experiment?", "limitations": "The authors have adequately addressed the limitation and potential negative social impact of their work.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The manuscript has the following strengths:\n1. The motivation is reasonable, and the authors proposed a framework to address the motivation.\n2. The experiments are solid, the authors provide comparible or better results in comparison with several baselines designated to handle long sequences.\n3. For such a complicated framework, the authors are able to illustrate clearly.\n\nThe manuscript has the following weaknesses:\n1. The proposed framework involves training of several models, including a memory model $g_\\theta$, a predictor function $q_\\phi$ and a uncertainty detector model $d_\\psi$. The whole model is too complicated in my point of view. \n\nSome typos:\n1. Title for Figure 2: \"final performanse\" should be \"final performance\"", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"ethics_review_area": ["I don’t know"]}, "extra_text": {}}]

Author rebuttal:

Decision: Accept