Backward pass

Author: MaximilianSchreff

scripts/nn/layers/embedding.dml

@@ -36,19 +36,50 @@ forward = function(matrix[double] indices, matrix[double] embedding_dict)
    *               embedding_dict[indices[i]].
    */
   n = nrow(indices)
-  d = ncol(embedding_dict)
+  v = nrow(embedding_dict)
 
   # Construct permutation-like matrix (one '1' per row, rest '0')
-  permutation = matrix(0, rows=n, cols=d)
+  permutation = matrix(0, rows=n, cols=v)
   for (i in 1:n) {
     permutation[i, as.integer(as.scalar(indices[i]))] = 1
   }
 
   embeddings = permutation %*% embedding_dict
 }
 
-backward = function()
-    return () {}
+backward = function(matrix[double] dout, matrix[double] indices, int v,
+      int padding_idx = -1)
+    return (matrix[double] dembedding_dict) {
+  /*
+   * Backward pass of embedding layer computes the gradients of the embedding
+   * dictionary.
+   *
+   * Inputs:
+   * - dout: Gradient of the output.
+   * - indices: Indices referring to embedding vectors of embedding dictionary
+   *            of shape n x 1 with each value in {1, ..., v}.
+   * - v: Embedding dictionary size.
+   * - padding_idx: Index of embedding vector of embedding dictionary which
+   *                should not be updated (i.e. gradients are 0). Use -1 if
+   *                there is no padding vector.
+   *
+   * Outputs:
+   * - dembedding_dict: Gradients of the dictionary of embedding vectors of 
+   *                    shape v x d.
+   */
+  n = nrow(indices)
+
+  # Construct permutation-like matrix (one '1' per row, rest '0')
+  permutation = matrix(0, rows=n, cols=v)
+  for (i in 1:n) {
+    permutation[i, as.integer(as.scalar(indices[i]))] = 1
+  }
+
+  dembedding_dict = t(permutation) %*% dout
+  if (padding_idx != -1) {
+    dembedding_dict[padding_idx] = matrix(0, rows=1, cols=ncol(dout))
+  }
+}
 
 init = function(int v, int d, int seed  = -1)
     return (matrix[double] embedding_dict) {
@@ -58,7 +89,7 @@ init = function(int v, int d, int seed  = -1)
    * Inputs:
    * - v: Embedding dictionary size.
    * - d: Embedding vector dimension.
-   * - seed: Optional random generation seed.
+   * - seed: Random generation seed.
    *
    * Output:
    * - embedding_dict: Embedding dictionary matrix of shape v x d.