scripts/nn/layers/gelu.dml

#-------------------------------------------------------------
#
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
#
#-------------------------------------------------------------

/*
 * Gaussian Error Linear Unit (GELU) nonlinearity layer.
 */

source("nn/layers/tanh.dml") as tanh

forward = function(matrix[double] X)
    return (matrix[double] out) {
  /*
   * Computes the forward pass for a GELU nonlinearity layer, via 
   * its tanh approximation.
   *
   * Performs an element-wise evaluation of 
   *    `GELU(x) = x * CDF(x)`.
   * where CDF is the cumulative distribution function of the 
   * standard normal distribution:
   *    `CDF(x) = 0.5 * (1 + erf(x/sqrt(2)))`
   * This implementation uses the tanh approximation:
   *    `CDF(x) =~ 0.5 * (1 + tanh(sqrt(2/pi) * (x + 0.044715x^3)))`
   *
   * Inputs:
   *  - X: Inputs, of shape (any, any).
   *
   * Outputs:
   *  - out: Outputs, of same shape as `X`.
   */
  cdf = 0.5 * (1 + tanh(sqrt(2 / pi) * (X + 0.044715 * X^3)))
  out = cdf * X
}

backward = function(matrix[double] dout, matrix[double] X)
    return (matrix[double] dX) {
  /*
   * Computes the backward pass for a GELU nonlinearity layer, via 
   * its tanh approximation.
   *
   * Inputs:
   *  - dout: Gradient wrt `out` from upstream, of same shape as `X`.
   *  - X: Previous input data matrix, of shape (any, any).
   *
   * Outputs:
   *  - dX: Gradient wrt `X`, of same shape as `X`.
   */
   a = sqrt(2 / pi)
   b = 0.044715
   T = tanh(a * (X + b * X^3))
   dT = 1 - T^2
   dX = dout * (0.5 * (1 + T) + 0.5 * X * dT * a * (1 + 3 * b * X^2))
}