diff --git a/src/mlpack/methods/ann/layer/CMakeLists.txt b/src/mlpack/methods/ann/layer/CMakeLists.txt index a91f441681..575fc8a9be 100644 --- a/src/mlpack/methods/ann/layer/CMakeLists.txt +++ b/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 diff --git a/src/mlpack/methods/ann/layer/batch_norm.hpp b/src/mlpack/methods/ann/layer/batch_norm.hpp index e4371c2627..662619ef80 100644 --- a/src/mlpack/methods/ann/layer/batch_norm.hpp +++ b/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. diff --git a/src/mlpack/methods/ann/layer/batch_norm_impl.hpp b/src/mlpack/methods/ann/layer/batch_norm_impl.hpp index 7abfe17233..b55bef3700 100644 --- a/src/mlpack/methods/ann/layer/batch_norm_impl.hpp +++ b/src/mlpack/methods/ann/layer/batch_norm_impl.hpp @@ -99,7 +99,7 @@ void BatchNorm::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::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 diff --git a/src/mlpack/methods/ann/layer/layer.hpp b/src/mlpack/methods/ann/layer/layer.hpp index 5a0c53764e..b4764d975f 100644 --- a/src/mlpack/methods/ann/layer/layer.hpp +++ b/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" diff --git a/src/mlpack/methods/ann/layer/layer_norm.hpp b/src/mlpack/methods/ann/layer/layer_norm.hpp new file mode 100644 index 0000000000..a978f2b022 --- /dev/null +++ b/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 + +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 + void Forward(const arma::Mat&& input, arma::Mat&& output); + + /** + * Backward pass through the layer. + * + * @param input The input activations. + * @param gy The backpropagated error. + * @param g The calculated gradient. + */ + template + void Backward(const arma::Mat&& input, + arma::Mat&& gy, + arma::Mat&& 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 + void Gradient(const arma::Mat&& input, + arma::Mat&& error, + arma::Mat&& 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 + 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 diff --git a/src/mlpack/methods/ann/layer/layer_norm_impl.hpp b/src/mlpack/methods/ann/layer/layer_norm_impl.hpp new file mode 100644 index 0000000000..9985102d10 --- /dev/null +++ b/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 +LayerNorm::LayerNorm() : + size(10), + eps(1e-8) +{ + // Nothing to do here. +} + +template +LayerNorm::LayerNorm( + const size_t size, const double eps) : + size(size), + eps(eps) +{ + weights.set_size(size + size, 1); +} + +template +void LayerNorm::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 +template +void LayerNorm::Forward( + const arma::Mat&& input, arma::Mat&& 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 +template +void LayerNorm::Backward( + const arma::Mat&& input, arma::Mat&& gy, arma::Mat&& 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 +template +void LayerNorm::Gradient( + const arma::Mat&& /* input */, + arma::Mat&& error, + arma::Mat&& 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 +template +void LayerNorm::serialize( + Archive& ar, const unsigned int /* version */) +{ + ar & BOOST_SERIALIZATION_NVP(gamma); + ar & BOOST_SERIALIZATION_NVP(beta); +} + +} // namespace ann +} // namespace mlpack + +#endif diff --git a/src/mlpack/methods/ann/layer/layer_types.hpp b/src/mlpack/methods/ann/layer/layer_types.hpp index d2aa77ceac..bc2f46dbe4 100644 --- a/src/mlpack/methods/ann/layer/layer_types.hpp +++ b/src/mlpack/methods/ann/layer/layer_types.hpp @@ -25,6 +25,7 @@ #include #include #include +#include #include #include #include @@ -45,10 +46,10 @@ namespace mlpack { namespace ann { - template class BatchNorm; template class DropConnect; template class Glimpse; +template class LayerNorm; template class Linear; template class LinearNoBias; template class LSTM; @@ -158,6 +159,7 @@ using LayerTypes = boost::variant< Glimpse*, HardTanH*, Join*, + LayerNorm*, LeakyReLU*, Linear*, LinearNoBias*, diff --git a/src/mlpack/tests/ann_layer_test.cpp b/src/mlpack/tests/ann_layer_test.cpp index 6a4e4d2b87..ed52dbd16b 100644 --- a/src/mlpack/tests/ann_layer_test.cpp +++ b/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, NguyenWidrowInitialization>(); + model->Predictors() = input; + model->Responses() = target; + model->Add >(); + model->Add >(256); + model->Add >(10, 2); + model->Add >(); + } + + ~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, NguyenWidrowInitialization>* model; + arma::mat input, target; + } function; + + BOOST_REQUIRE_LE(CheckGradient(function), 1e-4); +} + BOOST_AUTO_TEST_SUITE_END();