feat(a3c): add self-attention over sequences

 - common math considerations require knowledge of terms w.r.t other terms, and self-attention provides that.
 - we can also inspect it after training to determine if the attention is finding useful things
 - add dropout to LSTM layers

Author: justindujardin

libraries/mathy_python/mathy/agents/attention.py

@@ -0,0 +1,6 @@
+"""TensorFlow compatible import of SeqSelfAttention"""
+import os
+
+# required to tell library to use TF backend
+os.environ["TF_KERAS"] = "1"
+from keras_self_attention import SeqSelfAttention

libraries/mathy_python/mathy/agents/base_config.py

@@ -24,6 +24,9 @@ class Config:
     # (n - 1) previous timesteps
     prediction_window_size: int = 16
 
+    # Dropout to apply to LSTMs
+    dropout: float = 0.2
+
     # Whether to use the LSTM or non-reccurrent architecture
     use_lstm: bool = True
     units: int = 64
@@ -38,7 +41,7 @@ class Config:
     verbose: bool = False
     # Initial learning rate that decays over time.
     lr_initial: float = 0.01
-    lr_decay_steps: int = 1000
+    lr_decay_steps: int = 100
     lr_decay_rate: float = 0.96
     lr_decay_staircase: bool = True
     max_eps: int = 15000

libraries/mathy_python/mathy/agents/embedding.py

@@ -1,5 +1,4 @@
 from typing import Any, Dict, List, Optional, Tuple, Union
-
 import tensorflow as tf
 
 from mathy.agents.base_config import BaseConfig
@@ -11,41 +10,7 @@
 )
 from mathy.agents.densenet import DenseNetStack
 
-
-class BahdanauAttention(tf.keras.layers.Layer):
-    """Bahdanau Attention from:
-    https://www.tensorflow.org/tutorials/text/nmt_with_attention
-    
-    Used rather than the built-in tf.keras Attention because we want
-    to get the weights for visualization.
-    """
-
-    def __init__(self, units):
-        super(BahdanauAttention, self).__init__()
-        self.W1 = tf.keras.layers.Dense(units)
-        self.W2 = tf.keras.layers.Dense(units)
-        self.V = tf.keras.layers.Dense(1)
-
-    def call(self, query, values):
-        # query hidden state shape == (batch_size, hidden size)
-        # query_with_time_axis shape == (batch_size, 1, hidden size)
-        # values shape == (batch_size, max_len, hidden size)
-        # we are doing this to broadcast addition along the time axis to calculate the score
-        query_with_time_axis = tf.expand_dims(query, 1)
-
-        # score shape == (batch_size, max_length, 1)
-        # we get 1 at the last axis because we are applying score to self.V
-        # the shape of the tensor before applying self.V is (batch_size, max_length, units)
-        score = self.V(tf.nn.tanh(self.W1(query_with_time_axis) + self.W2(values)))
-
-        # attention_weights shape == (batch_size, max_length, 1)
-        attention_weights = tf.nn.softmax(score, axis=1)
-
-        # context_vector shape after sum == (batch_size, hidden_size)
-        context_vector = attention_weights * values
-        context_vector = tf.reduce_sum(context_vector, axis=1)
-
-        return context_vector, attention_weights
+from .attention import SeqSelfAttention
 
 
 class MathyEmbedding(tf.keras.Model):
@@ -95,6 +60,7 @@ def __init__(self, config: BaseConfig, **kwargs):
                 name="nodes_lstm",
                 time_major=False,
                 return_sequences=True,
+                dropout=self.config.dropout,
             ),
             merge_mode="sum",
         )
@@ -103,8 +69,13 @@ def __init__(self, config: BaseConfig, **kwargs):
             name="time_lstm",
             time_major=True,
             return_sequences=True,
+            dropout=self.config.dropout,
+        )
+        self.lstm_attention = SeqSelfAttention(
+            attention_activation="sigmoid",
+            name="self_attention",
+            return_attention=True,
         )
-        self.lstm_attention = BahdanauAttention(self.config.lstm_units)
 
     def call(self, features: MathyInputsType, train: tf.Tensor = None) -> tf.Tensor:
         output = tf.concat(
@@ -124,5 +95,5 @@ def call(self, features: MathyInputsType, train: tf.Tensor = None) -> tf.Tensor:
         output = self.time_lstm_norm(output)
         output = self.lstm_nodes(output)
         output = self.nodes_lstm_norm(output)
-        output, attention = self.lstm_attention(output, output)
+        output, attention = self.lstm_attention(output)
         return self.out_dense_norm(self.output_dense(output)), attention

libraries/mathy_python/requirements.txt

@@ -9,3 +9,5 @@ wasabi
 gym<=0.12.5
 tensorflow_probability
 typing-extensions>=3.7.4.1
+
+keras-self-attention