Explain My Surprise: Learning Efficient Long-Term Memory by Predicting Uncertain Outcomes
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 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 MemUP 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 [1; 2; 3; 4; 5; 6; 7; 8; 9]) and Transformers [10; 11; 12; 13; 14; 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 a few bits of information but for a very long time [16].
Evolution of neural architectures tackling the problem of learning long-term dependencies follows the path of increasing size and complexity [6; 8; 17].
For example, Transformers[18] 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 in the past and agent’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 can be measured in thousands or millions of time steps, such solutions seem unrealistic given the current hardware capabilities.
We propose a training method that allows an agent to find and store dependencies between temporarily distant events without the need to process intermediate steps.
Our contributions are the following:
36th Conference on Neural Information Processing Systems (NeurIPS 2022).
• We propose a new formulation for the problem of memory learning as maximisation of conditional mutual information between memory state and future outcomes.
• We introduce an original method MemUP ( Mem ory for U ncertainty P rediction) 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 is able to learn long-term dependencies better than baselines while propagating gradients only through a small fractions of a sequence in one optimization step.
Notably, the applications of MemUP are successful in supervised learning tasks with different loss functions as well as in reinforcement learning tasks.
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 to explain out these events.
Below we formulate this idea in the framework of information theory.
Lets consider a sequence τ with inputs { x t } Tt =0 and targets { y t } Tt =0 .
A model is trained to predict targets given inputs.
An ideal memory could just store all previous inputs { x i } t i =
at each step t and predict y t from them, but it is too computationally expensive.
In practice a memory model g θ has to selectively store some aggregation of this sequence in a memory state m t =
g θ ( x t , m t − 1 ) .
Then it should be decided what to write and what to remove from the memory.
Estimating past information importance.
Consider an arbitrary time step k .
The past element x t has useful information about the target y k , if y k depends on x t given current x k , i.e. p ( y k | x k , x t ) ̸
= p
y k | x k ) .
In this case, we can say that there is a temporal dependency with the length of k − t steps, that starts at step t and ends at step k .
The strength of temporal dependency or importance of remembering x t can be measured as an amount of change in the distribution of y k given a knowledge of x t .
Multiplication of the both parts of the inequality above by p ( x t | x k ) and application of
KL-divergence as a measure of the difference between distributions leads to the conditional mutual information as a measure of x t importance:
[Equation]
where x k is taken from sequences stored in some dataset or generated by an 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 θ :
[Equation]
The sum in eq.
2 allows us to process each memory update in a 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 to 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 y k prediction given x k .
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 in 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 ∗ t that stores all possible information from the past is available.
Then the conditional mutual information at a step t :
[Equation]
shows the potential utility of memory for prediction at this step.
When the number of elements in the sum in eq.
2 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 I ( y t , m ∗ t | x t ) directly, because m ∗ t is not available to us.
As seen from the eq.3, m ∗ t 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 entropy H ( y t | x t ) 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 I ( y t , m ∗ t | x t ) .
Specifically, for every pair of elements i , j if H ( y i | x i ) >
H ( y j | x j ) then I ( y i , m ∗ i | x i ) >
I ( y j , m ∗ j | x j ) .
This becomes true under one of two possible conditions: (1) H ( y t | x t ) ≫ H ( y t | m ∗ t , x t ) or (2)
H ( y t | m ∗ t , x t ) ∼ ϵ .
The first condition means that given a perfect memory and a model, one could make high-quality predictions for the task at hand.
The second condition means that the real distribution of the target variable y t has a similar amount of noise at each step.
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" H ( y t | x t ) .
We should note that estimation of I ( y j , m ∗ j | x j ) 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 a new objective function that allows us to train it only for steps where the potential gain from the past information is maximal:
[Equation]
here U t denotes the set of top- K indices of steps from t to T with the highest estimated local entropy
H ( y k | x k ) .
The hyperparameter K controls fraction of a sequence to be processed in one gradient step.
On the other hand, maximizing the mutual information between the memory m t and an arbitrary distant event at step k > t allows to learn long-term dependencies.
Optimization.
Directly optimizing mutual information can be a 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-Entropy (CE) between the empirical and model distribution for all selected high entropy events:
[Equation]
where m θt is a shortcut for m t
= g θ ( x t , m t − 1 ) , q ϕ is a predictor network with parameters ϕ , that estimates probability of y k given x k and m t .
For a detailed derivation of eq.
5 see Supplementary
Materials Section A.
Collection of elements with high uncertainty U t requires a special uncertainty detector d ψ model.
The main requirement for d ψ is to produce uncertainty estimates s t
= d ψ ( x t ) .
In the simplest case, the detector estimates “surprise” − log p ( y t | x t ) (prediction error), which can be seen as a single point estimate of uncertainty.
Other models that directly estimate the uncertainty or a whole distribution [23; 24] can also be used.
MemUP for Recurrent Neural Architectures
Recurrent neural net implementation of MemUP consists of (1) training of uncertainty detector model d ψ , (2) selection of elements with highest uncertainty U t and (3) training memory g θ and predictor q ϕ models.
A pseudocode for RNN MemUP training is shown in Supplementary Materials E.
Uncertainty detector training.
Information maximization reasoning does not limit the choice of solutions for uncertainty detector
d ψ .
In our experiments on algorithmic tasks we use a recurrent neural classifier trained with Cross-Entropy Loss.
In this case detector’s "surprise" − log d ψ ( y t | x t ) is used as an uncertainty estimate.
[Figure]
In Reinforcement Learning experiments we use distributional algorithm QRDQN [24] trained to predict discounted future returns R t ( y t =
R t in RL experiments).
As QRDQN approximates a whole distribution of the future returns, estimating uncertainty becomes a simple task.
Processing a sequence with the memory network.
MemUP allows to learn long term-dependencies without propagating gradients through the whole sequence.
To demonstrate that in our experiments we use Truncated BPTT to train memory network g θ .
In the majority of supervised learning experiments we set rollout length r between 10 and 60 steps.
In all Reinforcement Learning experiments r = 1 , i.e. recurrent memory is trained without actually using backpropagation through time.
In All Experiments g θ is a simple LSTM network with additional input encoder.
Selecting states with high uncertainty.
Given uncertainty estimates { s i } Ti = t for each future element in the sequence, we use softmax sampling p ( k ) =
e s k /τ / (cid:80)
Ti = t ( e s i /τ ) without replacement with τ = 0 .
To implement sampling without replacement Gumbel-Max Trick [25] is used.
A scheme illustrating a single gradient update is shown in Figure 1.
Related work
The problem of learning long-term dependencies has long been studied in both supervised learning [26; 27] and reinforcement learning domains [28].
However, most of the research related to recurrent networks architectures is aimed at solving the problems of exploding and vanishing gradients [2; 29; 7; 30].
The common trend is that many of the existing solutions to deal with vanishing gradients simultaneously increase the cost associated with the calculation of these gradients [29; 7].
For example, Transformers[10] generally require to store O ( N 2 ) of intermediate results for training on sequences of length N .
The recently proposed Linformer [31] 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 in combination with modern reinforcement learning algorithms [5; 32; 33; 34].
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 specifically modified DNC [7] architecture learned long-term dependencies of several hundred steps in the setting of 2D and 3D navigation tasks.
The downside of this architecture is that the agent needs access to the information about its absolute or relative location in a 2D/3D environment.
In another work, Ha and Schmidhuber [36] used a complex procedure with pre-training of memory and state embeddings.
Their overall learning procedure is similar to 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 al.
[17] demonstrated the state of the art results with the MERLIN algorithm.
MERLIN’s memory kept observation embeddings trained with a complex variational autoencoder in a DNC- like [7] list.
The memory module was also trained separately from the policy.
Hung et al.
[37] followed up on the MERLIN architecture’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 a 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 al.
[18] 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.
[16] proposed AMRL architecture: a simple modification to the LSTM architecture to combat the problem of learning long-term dependencies in environments with a significant amount of observational noise.
Supervised Learning Experiments
For evaluation and comparison of our method we use four tasks: Copy [40], Scattered copy, Add [2] and permuted sequential MNIST (pMNIST) [41].
All these tasks are benchmarks that are used for testing models with long-term memory.
In the original Copy task a sequence of length l ( l = 10 in out experiments) should be copied after a T − l steps ( T ∈ { 120 , 520 , 1020 , 5020 } in our experiments), and in Scattered copy task a model has to make predictions in locations inside range [ l, 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 10 K and 1 K, for pMNIST 60 K and 10 K.
For detailed tasks description see Section B.
In these experiments predictor q ϕ consists of a three-layered MLP and a recurrent input encoder.
The encoder E rnn consists of a single fully-connected layer followed by two LSTM-layers withhidden units and dropout probability 0 .
E rnn is used to encode bigger local context than original inputs corresponding to indices in the U t set
(it is also possible to use a feedforward network for this task).
Memory g θ has the same architecture as E rnn with separate weights.
We do not train a separate detector d ψ .
Uncertainty at each step is estimated by prediction error from the predictor module q ϕ .
For more information you can look at our implementation ( see section E).
In the Add Task the model is trained with MSE loss, which is a special case of CE loss under the model-assumption that target distribution is a Gaussian with unit variance.
Both memory and encoder are trained with the same Truncated BPTT length.
Truncation length equals 10 in all tasks with sequence length < 1000 , 20 in tasks with length ≥ 1000 and 30 , 60 in pMNIST 784 ,respectively.
We compare MemUP with LSTM [2] and Transformer [10] models from pytorch library, as well as a 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 E rnn with
layered
MLP on top.
In the case of SRNN baselines, we have replaced all LSTM layers with SRNN layers.
[Figure]
The Transformer model consists of a single FC layer followed by 4 self-attention layers and the same 3 layered MLP 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 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 a good quality of prediction in all tasks.
In truncated BPTT setting (solid lines)
MemUP is significantly better than LSTM tr, SRNN tr or
Transformer tr due to MemUP 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.
Moreover, 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 a hint at the starting position and location of a reward at the exit (see Figure 3.1).
The main advantage of this problem is that it allows to test the agent’s memory mechanisms in 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 [42; 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 a room filled with acid (see Figure 3.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 a +1 reward, while others take away health 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.3).
The episode ends when the agent’s health drops to zero or after 1050 time steps.
The most straightforward extension of the MemUP algorithm to a reinforcement learning problem is to use memory state m t as an additional input to an RL agent.
There are several possible solutions to combine MemUp’s training with agent’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 a single into an end-to-end network.
In these experiments, we test the simplest version: train MemUP 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 (see
Section E).
Inputs x t are formed from observations o t and previous action a t − 1 .
This time Observations are encoded using convolutional networks or fully connected networks.
We use discounted future returns
[Figure] [Figure]
R t = (cid:80) T i =
t γ i − t r i .
as targets y t .
In the experiments, MemUP is compared with the following baselines:
PPO-LSTM is a recurrent
version
of PPO [45]
with
a single LSTM-layer
We use
PPO
-LSTM
implementation from
the
RLPyt
library [46]
IMPALA-ST is an IMPALA [
] agent using Stabilized Transformer[10] architecture.
Stabilized
Transformer was presented by Parisotto et al.
For experiments, we use the only available open implementation of this algorithm [48].
AMRL is proposed by [16] for the Noisy T-Maze Task.
In the original study AMRL was compared with many different memory architectures including DNC, multi-layered LSTM and outperformed all of them on T-Maze-LN-100 task as well as other memory experiments.
AMRL is similar to
PPO-LSTM baseline, but extends LSTM with AMRL Layer.
In all experiments we use 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 g θ , PPO-LSTM and PPO-AMRL baselines.
Results for the T-Maze-LNR-1000 (10 times longer dependency than in [16]) environment are presented in Figure 4a.
The designation T-Maze-LNR-1000 means that the maze length (and corresponding tempral dependency) in each episode can be from 1000 tosteps.
Several rollout lengths are tested for RNN-based baselines.
We could not successfully train a single PPO-LSTM or AMRL agent even with TBPTT rollout of r = 1000 1 .
During MemUP pre-training phase, episodes are generated with a random policy.
We use the discounted future reward with γ = 0 as a prediction target.
A policy training for the MemUP agent starts after the memory pre-training phase is completed.
The time spent in the pre-training phase is marked with a dashed line in all figures.
As shown in Figure4a the agent learns almost instantly, given a memory 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, Section
D shows that MemUp can be trained to solve T-maze-LNR-20000.
Results for the Vizdoom-Two-Colors environment are shown in Figure 4b.
IMPALA-ST(attn=8) with a short attention span could not learn any reasonable policy.
On the other hand, IMPALAST(attn=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 is potentially enough to remember the color of the column or a previously collected items, PPO-LSTM and AMRL agents are unable to utilize long-term information from the past.
For MemUP we use the discounted future reward with γ = 0 .
8 as a prediction target.
The memory module g θ is trained with the rollout length of 1 step.
In Vizdoom-Two-Colors we set number of targets K = 3 .
Thus, to train the memory g θ we use 4 separate observations per episode for one gradient update.
In ViZDoom, we train PPO (MemUP+PPO) and QRDQN (MemUP+QRDQN) 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-training and evaluation.
As well, MemUP significantly outperforms baselines with comparable resource requirements, like IMPALA-ST (attn=8), PPOLSTM(r=8) and AMRL-MAX(r=8).
Sensitivity to the Noisy-TV problem
In section 2 we discussed conditions on true uncertainty of targets y t given a perfect memory m ∗ t .
The goal was to set conditions under which we can make a good guess about relative values of potential memory utility
I ( y t | m ∗ t , c t ) by estimating local entropy H ( y t | c t ) .
However, this may not be the case.
For example, some targets y t may be completely unpredictable and no information from the past can reduce the entropy for the distribution of y t : H ( y t | x t ) ∼ H ( y t | m ∗ t , x t ) .
Then the local entropy is large, while potential utility of memory for y t predictions is close to zero.
Training memory to predict such targets is pointless and prevents MemUP from focusing on events for which training long-term memory is essential.
The problem of over-emphasis on unpredictable surprising events is called
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 to the Noisy-TV problem, we have modified T-maze-LNR-100 environment.
Noisy T-maze With Distractors.
In this version, agent can receive +4 or -3 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 t [1]!
= 0 .
element in observation vector o t .
Decision points are distinguishable from each other.
Each decision point has a unique value of o t [1] .
Only in one decision point, the next reward depends on the agent’s action and the value of the hint ( o t [1] = 1 .
in this case ).
In other D decision points, the next rewards are completely random ( chosen with a 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.
[Figure]
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. 5) 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 U t set( by changing K ) improves the chance of selecting a future event that is a part of temporal dependency.
On the other hand, with an increase in U t , 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 in the number of distractions, memory learning speed decreases.
However, increasing the set U t 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 dependency fairly quickly even if the uncertainty detector select noisy events 80% of the time.
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, we 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.
Rnd-Pred: Same as MemUP, but steps for future outcome prediction are selected randomly and uniformly among all future timesteps.
This randomization verify whether predicting steps with high uncertainty helps MemUP to overcome the problem of learning long-term dependencies with small
BPTT rollouts.
Default: Same as MemUP, but it is trained to predict return R t at each step t , as 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 5).
But in this study we test this baseline on all intermediate Truncated BPTT lengths.
All versions have the same architecture, and hyperparametes.
They were trained in T-Maze-LNRand T-Maze-LNR-1000 problems.
Episodes were generated by a random strategy.
All versions are trained with 5 different Truncated BPTT lengths: 1, 5, 10, 50, 100.
Testing metric is a root-meansquared 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.e. in the order of increasing task complexity.
For every rollout length there is a corresponding absolute value of error between predictions of the model and targets
[Figure]
at the end of training ( 3 e 4 updates for T-Maze-LNR-100 and 3 e 5 for T-Maze-LNR-1000).
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 as the results of the PPO-LSTM agent on the T-Maze-LNR 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 a new method for training long-term memory.
The main idea is to train a memory network to predict future outcomes of high uncertainty and skip all others.
Predicting a small number of arbitrarily distant future outcomes substantially saves computational resources required to 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 significantly longer than the number of steps processed for a single gradient update.
None of the baseline architectures trained in a 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-Colors (Stabilized Transformer baseline has quadratic complexity while processingtimes more steps).
Even using fewer resources, MemUP outperforms non-truncated baselines that has simultaneous access to all elements of a sequence (including Transformer) on pMNIST(3136), Copy and Scattered
Copy tasks.
On T-maze and Vizdoom-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.
Acknowledgments and Disclosure of Funding
This work was partially supported by a grant for research centers in the field of artificial intelligence, provided by the Analytical Center for the Government of the Russian Federation in accordance with the subsidy agreement (agreement identifier 000000D730321P5Q0002) and the agreement with the
{+Moscow Institute of Physics and Technology dated November 1, 
+}
No. 70-2021-00138.
Checklist
For all authors.. (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] Our experimental results and ablation study align with our claims.
See Sections 6 and 8.
(b) Did you describe the limitations of your work?
[Yes] See Section 2 and Noisy-TV experiments(c)
Did you discuss any potential negative societal impacts of your work?
[No] No, Negative consequences of practical application of our idea (that we could think of), can be coupled with any improvement in DL/RL methods rather than specifically with our proposal.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] We describe our assumptions in Section 2 (L.83-96) (b) Did you include complete proofs of all theoretical results?
[Yes] See Section A. 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] Link to the code is given in Section E (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] We have included these details in Supplementary Materials (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] All graphs include mean and standard deviation computed over multiple runs with different random seed.
(d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] See Supplementary Materials .
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] We use RLPyt library and open implementation of Stabilized Transformer and SRNN for our baselines and cite authors in sections 5 and(b)
Did you mention the license of the assets?
[No] (c) Did you include any new assets either in the supplemental material or as a URL?
[No] (d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
[No] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[No]
A Mutual Information Maximization
Estimating mutual information
I ( y k , m θt | x k ) can be difficult task from a computational point of view.
Therefore, to make the task more trackable we use a lower bound for mutual information derived by [22].
Adapting it for the MemUP case, we get the following inequality:
[Equation]
where m θt is a shortcut notation denoting dependence of the memory state m t on the parameters of memory network g θ .
The distribution p is the true conditional distribution of future outcomes y k , which is unknown to us, but we can sample from it using training data.
The distribution q ϕ is approximated by the predictor network with parameters ϕ .
Inequality 6 follows from the relation between KL-divergence, cross-entropy (CE) and entropy: D KL ( p || q ϕ ) =
CE ( p, q ϕ ) − H ( p ) .
Using non-negativity property of KL-divergence leads to CE ( p, q ) ≥ H ( p ) .
The bound becomes exact when q ϕ ( y k | m θt , x k ) is equal to p ( y k | m θt , x k )
Since H ( y k | x k ) is independent from memory and predictor networks, maximizing the lower bound is the same as minimizing cross-entropy.
Therefore, the memory g θ and the predictor q ϕ can be jointly trained to maximize mutual information by simply minimizing cross entropy loss:
[Equation]
Specifically, the memory module g θ is trained by backpropagating gradients directly through the predictor network q ϕ .
B Supervised Learning Tasks
For accuracy evaluation we involve four tasks:
Copy, Scattered copy, Add and permuted MNIST (pMNIST).
All these tasks are benchmarks that are used for testing models with long-term memory.
Copy task.
This is an aligned sequence to sequence classification problem ( X → Y ).
Denote by
T length of X and Y .
The first 10 elements contain random uniform digits from set { 2 , . .
, 9 } that model has to copy and reproduce in the end of the sequence
where X i = 1 , i ∈ { T − 10 , . . .
, T − 1 } .
In all other positions except the first 10 and the last 10 X i = 1 , i ∈ { 10 , . . .
, T − 11 } .
Correspondingly
Y i = 0 , i ∈ { 0 , . . .
, T − 11 } except the last 10 where Y T −
+ i
=
X i , i ∈ { 0 , . . .
We split this task into four different sub-tasks by the length where T ∈ { 120 , 520 , 1020 , 5020 } .
Train dataset contains 10 k sequences and test consists of 1 k sequences.
Scattered copy task.
It is similar to the previous one.
The only difference is that locations where one has to make predictions are choosen at random in range [10 , . . .
, T − 1] .
In such locations
X i =
and Y
i equals to some element from the first 10 digits.
We split this task into four different sub-tasks by the length where T ∈ { 120 , 520 , 1020 , 5020 } .
This task is more complex because it requires locations detection and their ordered number counting.
As a metric in the previous two tasks we use negative log-likelihood (NLL).
Train dataset contains 10 k sequences and test consists of 1 k sequences.
Add task.
This is an aligned sequence to sequence regression problem ( X → Y ).
Each element in X is two dimensional.
The first dimension contains a random uniform value ∈ [0 , 1] and second component is 0 or 1 .
There are only three ones per sequence that correspond to two summands and their sum.
Values of Y elements equal to zero except one element at the sum location.
MSE metric was used here as a loss function and quality measure.
Train dataset contains 10 k sequences and test consists of 1 k sequences.
We split this task into four different sub-tasks by the length where
T ∈ { 100 , 500 , 1000 , 5000 } .
Permuted sequential MNIST task (pMNIST).
The dataset is obtained from the ordinary MNIST that includes 60 k train images and 10 k test images.
Each image is flattened to one dimensional vector.
[Table]
[Table]
It yields vectors of size 784 and we also add zero-value padding to obtain vectors of size 3136 .
A random permutation changes order of the vector elements consistently.
In training procedure we apply NLL loss function and measure quality by classification accuracy in %.
C Reinforcement Learning Tasks
Noisy T-Maze.
Namely, the agent starts at the very beginning of the central corridor next to the hint that it observes in the first step of the episode.
The agent can only move forward along the corridor.
At each step, the agent’s observations
o t are represented by a vector of length 3.
The first element o t [0] contains value of the hint: +1 or -1 at the first step, 0 otherwise.
The second element o t [1] is the indicator of reaching the intersection, which is 1 if the agent has reached the place of turn.
The last element o t [2] does not carry information and is a random noise, it is +1 or -1 with equal probability.
The agent receives the reward only at the end of the episode.
It is +4 if the correct turn was chosen and -3 otherwise.
The length of the maze is defined as the number of steps between the moment the agent sees the hint and the moment the agent has to make a choice based on the value of the hint.
The real maze length in experiments varies within 10 steps from episode to episode in order to decorrelate the observations of agents trained in parallel on multiple instances of environments (see PPO algorithms [45], A3C [5], IMPALA [47] ).
The naming of the environment indicates the minimum possible maze length, for example, the designation T-Maze-LNR-100 means that the maze length in each episode can be from 100 to 109 steps.
Vizdoom-Two-Colors.
Despite the fact that T-Maze environment allows simulating very long temporal dependencies, it is too simplistic.
To explore if MemUP augmentation can be scaled to much more complex tasks and environments it was tested in the Vizdoom-Two-Colors task introduced by Beeching et al.
In this task, the agent is placed in a room filled with acid (see Figure 3.2.) and constantly loses health.
Objects of green and red colors are scattered throughout the environment.
Items of one color replenish the agent’s health and give a +1 reward, while others take away health and give a -1 reward.
The correspondence between effects and items’ 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.
The episode ends when the agent’s health drops to zero or after
1050 time steps.
At each step agent receives a small "living reward" equal to +0.02.
Accordingly, the goal of the agent is to survive as long as possible in the environment by collecting items of the rewarding color.
To solve this problem, it is necessary to keep the color of the signal column in memory in order to be able to select objects of the correct color, even when the column is out of sight.
However, in the course of preliminary experiments, it turned out that a reactive agent without memory is able to learn a strategy in which it will always keep the column in sight, even if the room is filled with walls blocking the view from many angles.
Therefore, we created a new version of the environment, where the column disappears after the 45th step (see Figure 3.3.) and the number of walls in the room is significantly reduced.
Thus, the subtask of memorization became harder in comparison with the original version, and the subtask of navigation in the environment was simplified.
It is also worth noting that the agent does not receive information about the current health or the rewards received, since these observations actually provide the same information as the color of the column.
D T-Maze Scaling Experiments.
In this section, we test MemUP’s scaling ability in the T-maze-LNR environment.
For each version of the environment (with lengths 500, 1000, 5000, 10000, 20000) we train PPO+MemUP from scratch in 3 separate runs.
Picture 7a shows average episodic return after 500 training steps of PPO in the respective environment.
Picture 7b shows the root mean squared error of the MemUP memory and predictor modules after 200 training epochs.
The evaluation of the memory is conducted in the same way as in Section 8.
The memory module is trained with BPTT rollout r = 20 in all runs.
[Figure]
E Open Source Implementation
Implementation for MemUP can be found at: https://github.com/griver/memup
F Space Complexity for MemUP
Time complexity for MemUP would be the same as for RNNs trained with TBPTT.
In our method we focus on space complexity, as MemUP allows RNNs to learn long dependencies with short rollouts.
Space complexity for MemUP is O
( k rf + r ) , where:
• r - rollout length; • f - prediction frequency (you can predict from any intermediate step inside a single rollout); • k - number of prediction targets.
In all experiments (excluding Section 8) frequency (long-term predictions are made at the end of each rollout), therefore the space complexity is only O ( k + r ) .
We would like to note that we recommend using f < r in longer rollouts as it saves from the vanishing gradient problem inside a single rollout.
G Pseudocode for Memory Training
This section provides a simplified version of the memory training algorithm with feedforward predictor q ϕ .
Algorithm
Input: inputs
, targets
, d q
Train detector d on the sequence data
Get uncertainty estimates s for each step using d t = 0 and m =
None while t ≤ T do for i =
t to t + r do m =
g θ ( x i , m ) end for t =
min( T, t + r )
Using s [ t : T ] sample a set U t of K elements with highest uncertainty
Optimize CE loss from eq.
5: min θ,ϕ (cid:80) k
∈U t [ − log q ϕ ( y k | m, x k )]
stop_gradient( m ) end while