Merge pull request #1389 from ShikharJ/LayerNorm

Implement Layer Normalization.

src/mlpack/methods/ann/layer/CMakeLists.txt

@@ -41,6 +41,8 @@ set(SOURCES
   join.hpp
   join_impl.hpp
   layer.hpp
+  layer_norm.hpp
+  layer_norm_impl.hpp
   layer_traits.hpp
   layer_types.hpp
   leaky_relu.hpp

src/mlpack/methods/ann/layer/batch_norm.hpp

@@ -19,7 +19,7 @@ namespace mlpack {
 namespace ann /** Artificial Neural Network. */ {
 
 /**
- * Declaration of the Batch Normalization layer class. The layer tranforms
+ * Declaration of the Batch Normalization layer class. The layer transforms
  * the input data into zero mean and unit variance and then scales and shifts
  * the data by parameters, gamma and beta respectively. These parameters are
  * learnt by the network.

src/mlpack/methods/ann/layer/batch_norm_impl.hpp

@@ -99,7 +99,7 @@ void BatchNorm<InputDataType, OutputDataType>::Backward(
   const arma::mat inputMean = input.each_col() - mean;
   const arma::mat stdInv = 1.0 / arma::sqrt(variance + eps);
 
-  // Step 1: dl / dxhat *
+  // Step 1: dl / dxhat
   const arma::mat norm = gy.each_col() % gamma;
 
   // Step 2: sum dl / dxhat * (x - mu) * -0.5 * stdInv^3.
@@ -113,8 +113,8 @@ void BatchNorm<InputDataType, OutputDataType>::Backward(
 
   // Step 3: sum (dl / dxhat * -1 / stdInv) + variance *
   // (sum -2 * (x - mu)) / m.
-  g.each_col() += arma::sum(norm.each_col() % -stdInv, 1) + var %
-      arma::mean(-2 * inputMean, 1) / input.n_cols;
+  g.each_col() += (arma::sum(norm.each_col() % -stdInv, 1) + (var %
+      arma::mean(-2 * inputMean, 1))) / input.n_cols;
 }
 
 template<typename InputDataType, typename OutputDataType>

src/mlpack/methods/ann/layer/layer.hpp

@@ -19,6 +19,7 @@
 #include "convolution.hpp"
 #include "dropconnect.hpp"
 #include "glimpse.hpp"
+#include "layer_norm.hpp"
 #include "layer_types.hpp"
 #include "linear.hpp"
 #include "linear_no_bias.hpp"

src/mlpack/methods/ann/layer/layer_norm.hpp

@@ -0,0 +1,200 @@
+/**
+ * @file layer_norm.hpp
+ * @author Shikhar Jaiswal
+ *
+ * Definition of the Layer Normalization class.
+ *
+ * mlpack is free software; you may redistribute it and/or modify it under the
+ * terms of the 3-clause BSD license.  You should have received a copy of the
+ * 3-clause BSD license along with mlpack.  If not, see
+ * http://www.opensource.org/licenses/BSD-3-Clause for more information.
+ */
+#ifndef MLPACK_METHODS_ANN_LAYER_LAYERNORM_HPP
+#define MLPACK_METHODS_ANN_LAYER_LAYERNORM_HPP
+
+#include <mlpack/prereqs.hpp>
+
+namespace mlpack {
+namespace ann /** Artificial Neural Network. */ {
+
+/**
+ * Declaration of the Layer Normalization class. The layer transforms
+ * the input data into zero mean and unit variance and then scales and shifts
+ * the data by parameters, gamma and beta respectively over a single training
+ * data. These parameters are learnt by the network. Layer Normalization is
+ * different from Batch Normalization in the way that normalization is done
+ * for individual training cases, and the mean and standard deviations are
+ * computed across the layer dimensions, as opposed to across the batch.
+ *
+ * For more information, refer to the following papers,
+ *
+ * @code
+ * @article{Ba16,
+ *   author    = {Jimmy Lei Ba, Jamie Ryan Kiros and Geoffrey E. Hinton},
+ *   title     = {Layer Normalization},
+ *   volume    = {abs/1607.06450},
+ *   year      = {2016},
+ *   url       = {http://arxiv.org/abs/1607.06450},
+ *   eprint    = {1607.06450},
+ * }
+ * @endcode
+ *
+ * @code
+ * @article{Ioffe15,
+ *   author    = {Sergey Ioffe and
+ *                Christian Szegedy},
+ *   title     = {Batch Normalization: Accelerating Deep Network Training by
+ *                Reducing Internal Covariate Shift},
+ *   journal   = {CoRR},
+ *   volume    = {abs/1502.03167},
+ *   year      = {2015},
+ *   url       = {http://arxiv.org/abs/1502.03167},
+ *   eprint    = {1502.03167},
+ * }
+ * @endcode
+ *
+ * @tparam InputDataType Type of the input data (arma::colvec, arma::mat,
+ *         arma::sp_mat or arma::cube).
+ * @tparam OutputDataType Type of the output data (arma::colvec, arma::mat,
+ *         arma::sp_mat or arma::cube).
+ */
+template <
+  typename InputDataType = arma::mat,
+  typename OutputDataType = arma::mat
+>
+class LayerNorm
+{
+ public:
+  //! Create the LayerNorm object.
+  LayerNorm();
+
+  /**
+  * Create the LayerNorm object for a specified number of input units.
+  *
+  * @param size The number of input units.
+  * @param eps The epsilon added to variance to ensure numerical stability.
+  */
+  LayerNorm(const size_t size, const double eps = 1e-8);
+
+  /**
+   * Reset the layer parameters.
+   */
+  void Reset();
+
+  /**
+   * Forward pass of Layer Normalization. Transforms the input data
+   * into zero mean and unit variance, scales the data by a factor gamma and
+   * shifts it by beta.
+   *
+   * @param input Input data for the layer.
+   * @param output Resulting output activations.
+   */
+  template<typename eT>
+  void Forward(const arma::Mat<eT>&& input, arma::Mat<eT>&& output);
+
+  /**
+   * Backward pass through the layer.
+   *
+   * @param input The input activations.
+   * @param gy The backpropagated error.
+   * @param g The calculated gradient.
+   */
+  template<typename eT>
+  void Backward(const arma::Mat<eT>&& input,
+                arma::Mat<eT>&& gy,
+                arma::Mat<eT>&& g);
+
+  /**
+   * Calculate the gradient using the output delta and the input activations.
+   *
+   * @param input The input activations.
+   * @param error The calculated error.
+   * @param gradient The calculated gradient.
+   */
+  template<typename eT>
+  void Gradient(const arma::Mat<eT>&& input,
+                arma::Mat<eT>&& error,
+                arma::Mat<eT>&& gradient);
+
+  //! Get the parameters.
+  OutputDataType const& Parameters() const { return weights; }
+  //! Modify the parameters.
+  OutputDataType& Parameters() { return weights; }
+
+  //! Get the input parameter.
+  InputDataType const& InputParameter() const { return inputParameter; }
+  //! Modify the input parameter.
+  InputDataType& InputParameter() { return inputParameter; }
+
+  //! Get the output parameter.
+  OutputDataType const& OutputParameter() const { return outputParameter; }
+  //! Modify the output parameter.
+  OutputDataType& OutputParameter() { return outputParameter; }
+
+  //! Get the delta.
+  OutputDataType const& Delta() const { return delta; }
+  //! Modify the delta.
+  OutputDataType& Delta() { return delta; }
+
+  //! Get the gradient.
+  OutputDataType const& Gradient() const { return gradient; }
+  //! Modify the gradient.
+  OutputDataType& Gradient() { return gradient; }
+
+  //! Get the mean across single training data.
+  OutputDataType Mean() { return mean; }
+
+  //! Get the variance across single training data.
+  OutputDataType Variance() { return variance; }
+
+  /**
+   * Serialize the layer.
+   */
+  template<typename Archive>
+  void serialize(Archive& ar, const unsigned int /* version */);
+
+ private:
+  //! Locally-stored number of input units.
+  size_t size;
+
+  //! Locally-stored epsilon value.
+  double eps;
+
+  //! Locally-stored scale parameter.
+  OutputDataType gamma;
+
+  //! Locally-stored shift parameter.
+  OutputDataType beta;
+
+  //! Locally-stored parameters.
+  OutputDataType weights;
+
+  //! Locally-stored mean object.
+  OutputDataType mean;
+
+  //! Locally-stored variance object.
+  OutputDataType variance;
+
+  //! Locally-stored gradient object.
+  OutputDataType gradient;
+
+  //! Locally-stored delta object.
+  OutputDataType delta;
+
+  //! Locally-stored input parameter object.
+  InputDataType inputParameter;
+
+  //! Locally-stored output parameter object.
+  OutputDataType outputParameter;
+
+  //! Locally-stored normalized input.
+  OutputDataType normalized;
+}; // class LayerNorm
+
+} // namespace ann
+} // namespace mlpack
+
+// Include the implementation.
+#include "layer_norm_impl.hpp"
+
+#endif

src/mlpack/methods/ann/layer/layer_norm_impl.hpp

@@ -0,0 +1,124 @@
+/**
+ * @file layer_norm_impl.hpp
+ * @author Shikhar Jaiswal
+ *
+ * Implementation of the Layer Normalization class.
+ *
+ * mlpack is free software; you may redistribute it and/or modify it under the
+ * terms of the 3-clause BSD license.  You should have received a copy of the
+ * 3-clause BSD license along with mlpack.  If not, see
+ * http://www.opensource.org/licenses/BSD-3-Clause for more information.
+ */
+
+#ifndef MLPACK_METHODS_ANN_LAYER_LAYERNORM_IMPL_HPP
+#define MLPACK_METHODS_ANN_LAYER_LAYERNORM_IMPL_HPP
+
+// In case it is not included.
+#include "layer_norm.hpp"
+
+namespace mlpack {
+namespace ann { /** Artificial Neural Network. */
+
+template<typename InputDataType, typename OutputDataType>
+LayerNorm<InputDataType, OutputDataType>::LayerNorm() :
+    size(10),
+    eps(1e-8)
+{
+  // Nothing to do here.
+}
+
+template <typename InputDataType, typename OutputDataType>
+LayerNorm<InputDataType, OutputDataType>::LayerNorm(
+    const size_t size, const double eps) :
+    size(size),
+    eps(eps)
+{
+  weights.set_size(size + size, 1);
+}
+
+template<typename InputDataType, typename OutputDataType>
+void LayerNorm<InputDataType, OutputDataType>::Reset()
+{
+  gamma = arma::mat(weights.memptr(), 1, size, false, false);
+  beta = arma::mat(weights.memptr() + gamma.n_elem, 1, size, false, false);
+  gamma.fill(1.0);
+  beta.fill(0.0);
+}
+
+template<typename InputDataType, typename OutputDataType>
+template<typename eT>
+void LayerNorm<InputDataType, OutputDataType>::Forward(
+    const arma::Mat<eT>&& input, arma::Mat<eT>&& output)
+{
+  mean = arma::mean(input, 0);
+  variance = arma::var(input, 1, 0);
+
+  // Normalize the input.
+  output = input.each_row() - mean;
+  output.each_row() /= arma::sqrt(variance + eps);
+
+  // Reused in the backward and gradient step.
+  normalized = output;
+
+  // Scale and shift the output.
+  output.each_row() %= gamma;
+  output.each_row() += beta;
+}
+
+template<typename InputDataType, typename OutputDataType>
+template<typename eT>
+void LayerNorm<InputDataType, OutputDataType>::Backward(
+    const arma::Mat<eT>&& input, arma::Mat<eT>&& gy, arma::Mat<eT>&& g)
+{
+  const arma::mat inputMean = input.each_row() - mean;
+  const arma::mat stdInv = 1.0 / arma::sqrt(variance + eps);
+
+  // dl / dxhat
+  const arma::mat norm = gy.each_row() % gamma;
+
+  // sum dl / dxhat * (x - mu) * -0.5 * stdInv^3.
+  const arma::mat var = arma::sum(norm % inputMean, 0) %
+      arma::pow(stdInv, 3.0) * -0.5;
+
+  // dl / dxhat * 1 / stdInv + variance * 2 * (x - mu) / m +
+  // dl / dmu * 1 / m.
+  g = (norm.each_row() % stdInv) + (inputMean.each_row() %
+      var * 2 / input.n_rows);
+
+  // sum (dl / dxhat * -1 / stdInv) + variance *
+  // (sum -2 * (x - mu)) / m.
+  g.each_row() += (arma::sum(norm.each_row() % -stdInv, 0) + (var %
+      arma::mean(-2 * inputMean, 0))) / input.n_rows;
+}
+
+template<typename InputDataType, typename OutputDataType>
+template<typename eT>
+void LayerNorm<InputDataType, OutputDataType>::Gradient(
+    const arma::Mat<eT>&& /* input */,
+    arma::Mat<eT>&& error,
+    arma::Mat<eT>&& gradient)
+{
+  gradient.set_size(size + size, 1);
+
+  // Step 5: dl / dy * xhat.
+  gradient.submat(0, 0, gamma.n_elem - 1, 0) = arma::reshape(arma::sum(
+      normalized % error, 0), normalized.n_cols, 1);
+
+  // Step 6: dl / dy.
+  gradient.submat(gamma.n_elem, 0, gradient.n_elem - 1, 0) =
+      arma::reshape(arma::sum(error, 0), error.n_cols, 1);
+}
+
+template<typename InputDataType, typename OutputDataType>
+template<typename Archive>
+void LayerNorm<InputDataType, OutputDataType>::serialize(
+    Archive& ar, const unsigned int /* version */)
+{
+  ar & BOOST_SERIALIZATION_NVP(gamma);
+  ar & BOOST_SERIALIZATION_NVP(beta);
+}
+
+} // namespace ann
+} // namespace mlpack
+
+#endif

src/mlpack/methods/ann/layer/layer_types.hpp

@@ -25,6 +25,7 @@
 #include <mlpack/methods/ann/layer/elu.hpp>
 #include <mlpack/methods/ann/layer/hard_tanh.hpp>
 #include <mlpack/methods/ann/layer/join.hpp>
+#include <mlpack/methods/ann/layer/layer_norm.hpp>
 #include <mlpack/methods/ann/layer/leaky_relu.hpp>
 #include <mlpack/methods/ann/layer/flexible_relu.hpp>
 #include <mlpack/methods/ann/layer/log_softmax.hpp>
@@ -45,10 +46,10 @@
 namespace mlpack {
 namespace ann {
 
-
 template<typename InputDataType, typename OutputDataType> class BatchNorm;
 template<typename InputDataType, typename OutputDataType> class DropConnect;
 template<typename InputDataType, typename OutputDataType> class Glimpse;
+template<typename InputDataType, typename OutputDataType> class LayerNorm;
 template<typename InputDataType, typename OutputDataType> class Linear;
 template<typename InputDataType, typename OutputDataType> class LinearNoBias;
 template<typename InputDataType, typename OutputDataType> class LSTM;
@@ -158,6 +159,7 @@ using LayerTypes = boost::variant<
     Glimpse<arma::mat, arma::mat>*,
     HardTanH<arma::mat, arma::mat>*,
     Join<arma::mat, arma::mat>*,
+    LayerNorm<arma::mat, arma::mat>*,
     LeakyReLU<arma::mat, arma::mat>*,
     Linear<arma::mat, arma::mat>*,
     LinearNoBias<arma::mat, arma::mat>*,

src/mlpack/tests/ann_layer_test.cpp

@@ -1396,7 +1396,7 @@ BOOST_AUTO_TEST_CASE(BatchNormTest)
 /**
  * BatchNorm layer numerically gradient test.
  */
-BOOST_AUTO_TEST_CASE(GradientBatchNormLayerTest)
+BOOST_AUTO_TEST_CASE(GradientBatchNormTest)
 {
   // Add function gradient instantiation.
   struct GradientFunction
@@ -1719,4 +1719,83 @@ BOOST_AUTO_TEST_CASE(GradientAtrousConvolutionLayerTest)
   BOOST_REQUIRE_LE(CheckGradient(function), 1e-3);
 }
 
+/**
+ * Tests the LayerNorm layer.
+ */
+BOOST_AUTO_TEST_CASE(LayerNormTest)
+{
+  arma::mat input, output;
+  input << 5.1 << 3.5 << arma::endr
+        << 4.9 << 3.0 << arma::endr
+        << 4.7 << 3.2 << arma::endr;
+
+  LayerNorm<> model(input.n_cols);
+  model.Reset();
+
+  model.Forward(std::move(input), std::move(output));
+  arma::mat result;
+  result << 1.2247 << 1.2978 << arma::endr
+         << 0 << -1.1355 << arma::endr
+         << -1.2247 << -0.1622 << arma::endr;
+
+  CheckMatrices(output, result, 1e-1);
+  result.clear();
+
+  output = model.Mean();
+  result << 4.9000 << 3.2333 << arma::endr;
+
+  CheckMatrices(output, result, 1e-1);
+  result.clear();
+
+  output = model.Variance();
+  result << 0.0267 << 0.0422 << arma::endr;
+
+  CheckMatrices(output, result, 1e-1);
+}
+
+/**
+ * LayerNorm layer numerically gradient test.
+ */
+BOOST_AUTO_TEST_CASE(GradientLayerNormTest)
+{
+  // Add function gradient instantiation.
+  struct GradientFunction
+  {
+    GradientFunction()
+    {
+      input = arma::randn(10, 256);
+      arma::mat target;
+      target.ones(1, 256);
+
+      model = new FFN<NegativeLogLikelihood<>, NguyenWidrowInitialization>();
+      model->Predictors() = input;
+      model->Responses() = target;
+      model->Add<IdentityLayer<> >();
+      model->Add<LayerNorm<> >(256);
+      model->Add<Linear<> >(10, 2);
+      model->Add<LogSoftMax<> >();
+    }
+
+    ~GradientFunction()
+    {
+      delete model;
+    }
+
+    double Gradient(arma::mat& gradient) const
+    {
+      arma::mat output;
+      double error = model->Evaluate(model->Parameters(), 0, 256, false);
+      model->Gradient(model->Parameters(), 0, gradient, 256);
+      return error;
+    }
+
+    arma::mat& Parameters() { return model->Parameters(); }
+
+    FFN<NegativeLogLikelihood<>, NguyenWidrowInitialization>* model;
+    arma::mat input, target;
+  } function;
+
+  BOOST_REQUIRE_LE(CheckGradient(function), 1e-4);
+}
+
 BOOST_AUTO_TEST_SUITE_END();