feat(a3c): add bahdanau attention layer

 - the keras layer won't give you its weights (greedy!)
 - from: https://www.tensorflow.org/tutorials/text/nmt_with_attention

Author: justindujardin

libraries/mathy_python/mathy/agents/base_config.py

@@ -36,7 +36,8 @@ class Config:
     init_model_from: Optional[str] = None
     train: bool = False
     verbose: bool = False
-    lr: float = 2e-4
+    # Initial learning rate that decays over time.
+    lr: float = 0.01
     max_eps: int = 15000
     # How often to write histograms to tensorboard (in training steps)
     summary_interval: int = 100

libraries/mathy_python/mathy/agents/embedding.py

@@ -12,6 +12,42 @@
 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
+
+
 class MathyEmbedding(tf.keras.Model):
     def __init__(self, config: BaseConfig, **kwargs):
         super(MathyEmbedding, self).__init__(**kwargs)
@@ -26,7 +62,9 @@ def __init__(self, config: BaseConfig, **kwargs):
         # +1 for the time - removed for ablation
         # +2 for the problem type hashes
         self.concat_size = 3
-        self.values_dense = tf.keras.layers.Dense(self.config.units, name="values_input")
+        self.values_dense = tf.keras.layers.Dense(
+            self.config.units, name="values_input"
+        )
         # self.time_dense = tf.keras.layers.Dense(self.config.units, name="time_input")
         self.type_dense = tf.keras.layers.Dense(self.config.units, name="type_input")
         self.in_dense = tf.keras.layers.Dense(
@@ -64,7 +102,7 @@ def __init__(self, config: BaseConfig, **kwargs):
             time_major=True,
             return_sequences=True,
         )
-        self.lstm_attention = tf.keras.layers.Attention()
+        self.lstm_attention = BahdanauAttention(self.config.lstm_units)
 
     def call(self, features: MathyInputsType, train: tf.Tensor = None) -> tf.Tensor:
         output = tf.concat(
@@ -84,5 +122,5 @@ def call(self, features: MathyInputsType, train: tf.Tensor = None) -> tf.Tensor:
         output = self.time_lstm_norm(output)
         output, state_h, state_c = self.lstm_nodes(output)
         output = self.nodes_lstm_norm(output)
-        output = self.lstm_attention([output, state_h])
-        return self.out_dense_norm(self.output_dense(output))
+        output, attention = self.lstm_attention(output, output)
+        return self.out_dense_norm(self.output_dense(output)), attention

libraries/mathy_python/mathy/agents/policy_value_model.py

@@ -147,7 +147,7 @@ def call(
         start = time.time()
         inputs = features_window
         # Extract features into contextual inputs, sequence inputs.
-        sequence_inputs = self.embedding(inputs)
+        sequence_inputs, attention_weights = self.embedding(inputs)
         sequence_mean = tf.reduce_mean(sequence_inputs, axis=1)
         values = self.value_net(sequence_mean)
         reward_logits = self.reward_net(sequence_mean)
@@ -163,7 +163,7 @@ def call(
                     time.time() - start, batch_size
                 )
             )
-        return logits, values, mask_logits, reward_logits, grouping
+        return logits, values, mask_logits, reward_logits, attention_weights
 
     def apply_pi_mask(
         self, logits: tf.Tensor, features_window: MathyInputsType,