Why We Need Working Memory?

. . .

. . .

Recurrent Memory:

\(h_{10}\)

. . .

. . .

\(h_9\)

\(h_{19}\)

. . .

. . .

. . .

Problem with RNNs:

Solutions for Vanishing Gradients

Fight RNN problems by building more complex RNNs

Long-Short Term Memory: LSTM

Differential Neural Computer: DNC

Window Based Memory:

. . .

. . .

Information

Gradients

Soft-Attention:

. . .

. . .

\(e_{10}\)

\(e_{12}\)

Soft-Attention:

. . .

. . .

\(e_{10}\)

\(e_{12}\)

\(a_{10}\)

Embeddings

Query

\(a_{12}\)

Context Vector \(c_{20} = \sum_t a_t e_t\)

Attention weight: \(a_t = {e_{t}}^T q/\sum_i {e_{i}}^T q \)

\(q\)

Transformer Basics: Self-Attention

Transformer Basics: Self-Attention

Self-Attention computation:

Each \(z_t\) contains relevant information about \(o_t\) collected over all steps in

Memory Window:

Full Transformer

This is how real Transformer looks:

Transformer:

• All time-steps in a memory window attends to all other time-steps
• For each time-step you have: Query \(q_t\), Key \(k_t\), Value \(v_t\)
• Transformer has N attention heads!
• Transformer uses positional encoding to relate different time steps temporarily
• You repeat this process for several layers!

Transformer

Attention is All You Need!

Problems with RNN and Transformers

• Temporal dependency should fully fit into TBPTT or Attention Span to be learned

• You need to store all intermediate computation to implement backpropogation over Attention Span/TBPTT

Problems with RNN and Transformers 

• RNN-based solutions:

  • Vanishing/Exploding Gradients

  • Hidden State Bottleneck
  • Trajectory Noise Sensitivity
  • Linear complexity without attention

• Transformers-based solutions:

  • Quadratic space and time complexity for most architectures
  • Less stable in RL
  • Linear space and time complexity for some architectures (which are not equal to full Self-Attention)

BigBird (Zaheer et al, 2020)

TBPTT and Attention Spans in RL: 

• RNN and MANN:

  • 100-250 steps for AMRL

  • 40-80 steps for R2D2

  • Best results that i know 500+ steps with MERLIN (Wayne et al, 2018)

• Transformers:

  • 512 steps in Stabilized Transformer
  • 2048 for GPT3 in Supervised Learning
  • 4096 for Linformer and BigBird

Can we find temporal dependencies locally?

Information from red timestep could help at blue timestep.

Can we find temporal dependencies locally?

Information from red timestep could help at blue timestep.

Markov property doesn't work for these timesteps: 

How important it is to remember the red timestep?

\(f_t\) can be Q-value, \(s_{t+1}\), \(r_{t+1}\), etc.

Can we find temporal dependencies locally?

Information from red timestep could help at blue timestep.

How important it is to remember the red timestep?

Memory's Objective

The best memory would maximize the following sum: 

Memory's Objective

Train memory to maximize MI for all future steps: 

\(O(T^2)\) in time!

instead of optimizing the whole second sum we can focus only on terms with highest Mutual Information, i.e. optimize w.t.r to moments where memory

still requires to process full sequence for to update \(m_t\)

Instead of optimizing \(m_t\) with respect for the whole sum \(\sum^{T}_{i=t}\), we can select the most w

Observation: not all steps in sum \(\sum^{T}_{i=t}\) are equally valuable.

Memory's Objective

Train memory to maximize MI for all future steps: 

Memory's Objective

Train memory to maximize MI for all future steps: 

Finding Important Moments

Goal: Find moments where memory is the most important for model's predictions!

Finding Important Moments

Problems:

• can't find steps with highest MI without computing MI for all steps t and i
• can't estimate MI for "perfect memory" that we don't have

Intuition:

Two main steps:

• First we find the most surprising/uncertain events in a sequence
• Then we search sequence for information that explains these events and store it in memory

Idea

Resume

• Train memory to maximize MI at moments that would benefit the most from using memory:​

         where \(U_t\) is a set of moments that can benefit the most from memory; \(|U| \ll T\). 

• Find these moments by estimating local uncertainty \(H(f_t| o_t, a_t)\)

Practical Implementation for RNN

Learning is divided in two phases:

• Memory Pretraining Phase - we train memory without improving agent's policy
• Policy Learning Phase - We train PPO agent, that receives output of the pretrained memory as an additional input.

Practical Implementation for RNN

1. Get estimates of local uncertainty for each time step in the episode

Practical Implementation for RNN

1. Get estimates of local uncertainty for each time step in the episode
2. Select top K moments in episode with highest uncertainty

Practical Implementation for RNN

1. Get estimates of local uncertainty for each time step in the episode
2. Select top K moments in episode with highest uncertainty
3. Pass through the episode and train a memory model to store information that improves predictions at the selected moments

Uncertainty Detection Example:

Environment: Gym-Minigrid (Chevalier-Boisvert et al, 2018)

Experiments

Baselines

We compare MemUP (Memory via Uncertainty Prediction) with the following baselines:

• PPO-LSTM is a recurrent version of PPO with a single LSTM-layer.
• IMPALA-ST is an IMPALA agent with Transformer-based architecture. The Stabilized Transformer architecture was presented by Parisotto et al. (2020)
• AMRL is proposed by Beck et al. (2020) for the Noisy T-Maze Task.  AMRL
  is similar to PPO-LSTM baseline, but extends LSTM with AMRL Layer.

IMPALA-ST

Experiments: Noisy T-Maze

Noisy T-Maze Environment:

• Hard Task in terms of Memory
• Simple Task in every other aspect 

Env Details:

• This version was proposed in AMRL paper
• Noisy Observations
• Determine the long-term dependency between hint at the start and reward at the end
• Learn to remember the hint, to achieve maximum reward:
  • +4 for matching the hint, -3 otherwise 

Experiments: Noisy T-Maze

Noisy T-Maze-100

Noisy T-Maze-1000

• All algorithms process sequences shorter than the temporal dependency
• Only MemUp solves the problem consistently

Experiments: Vizdoom

Details:

• Agent walks on acid (-4 health)
• +25 health for objects with a right color, otherwise -25 health
• The column hints the right color of objects 
• The column disappears after 45 steps
• +1 reward for the right color, -1 otherwise

ViZDoom-two-colors:

• 3D environment
• Complex observation space
• ~100 steps long temporal dependency
• Proposed in Breeching et al (2019)

Experiments: Vizdoom

Uncertainty Detector Demonstration

Experiments: Vizdoom

Results

• MemUP solves Vizdoom-Two-Colors

Experiments: Vizdoom

Memory Demonstration

Experiments: Ablation

Ablation Baselines

Memory Baselines:

• MemUP - our proposed algorithm
• Rnd-Pred - Same as MemUP, but events for prediction are selected randomly and uniformly among all future timesteps.
• 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.

Experiment Setup:

• Noisy T-maze environment, but we ignore policy learning phase
• Baselines are compared in terms of their ability to predict \(f_t\)
• Separate test set of 100 episodes

Experiments: Ablation

Results

Request: Comparison with the Papers

• Go-Explore (https://arxiv.org/abs/1901.10995)
• Neural Turing Machines (https://arxiv.org/abs/1410.5401)
• Recurrent Independent Mechanisms (https://arxiv.org/abs/1909.10893)
• PlaNet (https://arxiv.org/abs/1811.04551)
• Big Bird (https://arxiv.org/abs/2007.14062)??

Go-Explore

Go-Explore improves exploration in environments with sparse rewards

Go-Explore vs MemUP

Similarities:

• Both select some "interesting" states and do something with them :)

Differences:

Go-Explore:

1. Goal: Exploration-Exploitation dilemma  
2. "Interesting" state criterion: Number of Visits / Total Reward
3. Interesting states utilization: Explore from them

MemUP:

1. Goal: Learning working memory
2. "Interesting" state criterion: Lack of information
3. Interesting states utilization: Find information about them

NTMs and RIMs

Recurrent neural networks:

Training RNN on Sequences:

LSTM

LSTM

LSTM

LSTM

LSTM

NTMs and RIMs

Neural Turing Machines:

NTMs and RIMs

Recurrent Independent Mechanisms:

NTMs and RIMs vs MemUP

Similarities:

• Same Goal: Learn and use memory on sequential tasks

Differences:

NTM and RIM:

1. Modify architecture
2. Main Idea: Keep memory unaltered to fight with Vanishing Gradients
3. Temporal dependency should fit in BPTT

MemUP:

1. Modifies training procedure
2. Main Idea: Predict long-term future selected by the lack of information
3. No need to fit temporal dependency in BPTT

PlaNet

PlaNet builds a good model of the environment then plans with it

Two Main PlaNet improvements:

1.

PlaNet

PlaNet builds a good model of the environment then plans with it

Two Main PlaNet improvements:

2.

PlaNet vs MemUP

Similarities:

• Learning a good model in POMDP includes learning a good memory
• Learns by predicting more than 1 step into future

Differences:

PlaNet's Model:

1. Latent overshooting: prediction for the next N steps
2. Generate for all intermediate steps in prediction horizon
3. Focus on short term predictions:
   1. plan 12 steps ahead,
   2. ?did not use Latent Overshooting?

MemUP:

1. Predict for arbitrary distant future events
2. No need to to make intermediate predictions
3. Focus on long-term memory: 100, 1000 steps, etc.

BigBird

• Big Bird has Linear complexity in space instead of Quadratic complexity
• Can Process Longer Sequences
• Needs more Layers

Big Bird vs MemUP

Similarities:

• We fight the same problem

Differences (more like parallels):

Big Bird:

1. Different Algorithms using different architectures
2. Adds Global Tokens attending to everything
3. Linear complexity in Space for each gradient update
4. Quadratic complexity for each sequence

MemUP:

1. Different Algorithms using different architectures
2. "Selects some of the tokens that attend to everything"
3. Constant complexity in Space for each gradient update
4. Linear complexity for each sequence

Further Research Prespectives 1/2

Problems with Current Implementation:

• Two-Phase Learning Process in RL
  • There are some concerns but we probably can train memory and policy in parallel
•  Different Prediction Targets
  • Try to predict next observation, next reward

Appilcation to Other Domains:

• Offline Reinforcement Learning
  • No need to fix two-phase learning
• Supervised Learning: NLP tasks in particular
  • Core Idea behind MemUP is independent of RL specific properties
  • No need to fix two-phase learning

Further Research Prespectives 2/2

Major Research Directions:

• Fight with Noisy-TV problem
  • Estimate Epistemic (reducible) and Aleatoric (irreducible) Entropy 
  • Estimating entropy reduction progress
•  Combining MemUP with Transformer Archtechures
  • Select which observations/tokens to store in transformer's attention window
  • Local context is processed by transformer while MemUP based memory process long-term dependencies
• Online Learning (only for RL applications)
  • Learn memory/model in parallel with policy
• Meta-Learning Applications
  • Adapt to new tasks by selecting and storing samples with maximal MI for this task