title: TensorFlow Neural Network LSTM: Loss Function for Temporal Error
body: <p>Let's say we are trying to predict future earthquakes based on previous earthquakes. A visualisation of part of a dataset containing sequences of earthquakes may look like this (features to the left, labels to the right, x=month, y=earthquake magnitude):
  <a href="https://i.stack.imgur.com/myRrj.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/myRrj.png" alt="features - labels"></a>
  You and I can see that the pattern here is that when there are two magnitude 5 earthquakes in succession, they are followed by a magnitude 8 earthquake then a magnitude 4 earthquake. Those earthquakes may appear with a variation of 3-5 months away, but the magnitudes are the same when they do.</p>
  
  <p>Training an LSTM on this data with a mean squared error loss function produces the following predictions (the dashed line):
  <a href="https://i.stack.imgur.com/IKYwH.png" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/IKYwH.png" alt="predictions"></a>
  It's done exactly as we've asked it, it's learned to predict the mean per timestep. However, this is a poor prediction in terms of what the model is intended to do.</p>
  
  <p>For you or I looking at the data, we could surmise that the first predicted earthquake is a magnitude 8, and taking into account the variance in timestep our best prediction would be that it happens at timestep 8. We'd essentially be doing a mean error in the temporal axis, as well as the value axis.</p>
  
  <p><strong>What loss function can be used to do this in TensorFlow?</strong></p>
  
  <p>I've been searching but only finding material relating to per timestep error.
  I started trying a few methods that I thought might give better results like <code>tf.sort</code> on the labels and predictions before performing mean squared error, and then also <code>tf.argsort</code> on the labels and predictions and doing mean squared error on them to give a temporal error, but argsort has no gradients for back-propagation so started trying to write custom gradients, which just feels like overkill when a solution must already exist? (p.s. I'm comfortable with TensorFlow but I'm no expert).</p>
  
  <p>If someone can point me in the right direction I'd be hugely appreciative!</p>
  
  <p>Here's the code that demonstrates the example above:</p>
  
  <pre><code>import tensorflow as tf
  import numpy as np
  import matplotlib.pyplot as plt
  
  features = np.array([[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], 
                       [0.0, 0.0, 0.5, 0.0, 0.5, 0.0, 0.0], 
                       [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], 
                       [0.0, 0.0, 0.5, 0.0, 0.5, 0.0, 0.0], 
                       [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], 
                       [0.0, 0.0, 0.5, 0.0, 0.5, 0.0, 0.0]]).reshape(6, 7, 1).astype(np.float32)
  labels = np.array([[0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
                     [0.8, 0.0, 0.0, 0.4, 0.0, 0.0, 0.0],
                     [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
                     [0.0, 0.0, 0.8, 0.0, 0.0, 0.4, 0.0],
                     [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
                     [0.0, 0.8, 0.0, 0.0, 0.4, 0.0, 0.0]]).reshape(6, 7, 1).astype(np.float32)
  
  plt.figure(figsize=(20, 2))
  plt.plot(np.concatenate((features[:, :, 0].T, labels[:, :, 0].T)) * 10)
  plt.axvline(features.shape[1]-1)
  plt.title('&lt;- Features vs. Labels -&gt;')
  plt.xlabel('Month')
  plt.ylabel('Magnitude')
  plt.show()
  
  def train_input_fn(features, labels):
      dataset = tf.data.Dataset.from_tensor_slices((features, labels))
      return dataset.shuffle(len(features)).repeat().batch(len(features))
  def predict_input_fn(features):
      dataset = tf.data.Dataset.from_tensor_slices(features)
      return dataset.batch(len(features))
  
  def model_fn(features, labels, mode):
      lstm = tf.keras.layers.Bidirectional(tf.keras.layers.LSTM(units=100, return_sequences=True, input_shape=(7, 1)))(features)
      outputs = tf.keras.layers.Dense(units=1, activation='linear')(lstm)
  
      if mode == tf.estimator.ModeKeys.PREDICT:
          return tf.estimator.EstimatorSpec(mode, predictions=outputs)
  
      loss = tf.losses.mean_squared_error(labels, outputs)
      optimiser = tf.train.AdamOptimizer(learning_rate=0.01)
      train_op = optimiser.minimize(loss, global_step=tf.train.get_global_step())
      return tf.estimator.EstimatorSpec(mode, loss=loss, train_op=train_op)
  
  estimator = tf.estimator.Estimator(model_fn=model_fn)
  estimator.train(input_fn=lambda:train_input_fn(features, labels), steps=100)
  predictions = np.array(list(estimator.predict(input_fn=lambda:predict_input_fn(features))))
  
  plt.figure(figsize=(20, 2))
  plt.plot(np.concatenate((features[:, :, 0].T, labels[:, :, 0].T)) * 10)
  plt.plot(np.concatenate((features[:, :, 0].T, predictions[:, :, 0].T)) * 10, ls='--')
  plt.axvline(features.shape[1]-1)
  plt.title('&lt;- Features vs. Labels -&gt; (with model predictions \'--\')')
  plt.xlabel('Month')
  plt.ylabel('Magnitude')
  plt.show()
  </code></pre>