From 45f162f1fc707af200ca5442bd73ca558bcfe1e9 Mon Sep 17 00:00:00 2001
From: nishantkr18 <nniishantkumar@gmail.com>
Date: Fri, 29 May 2020 20:00:24 +0530
Subject: [PATCH] Forward and backward update for NoisyLinear

---
 src/mlpack/methods/ann/layer/noisylinear.hpp  | 33 +++++++++-
 .../methods/ann/layer/noisylinear_impl.hpp    | 63 ++++++++++++++++---
 2 files changed, 87 insertions(+), 9 deletions(-)

diff --git a/src/mlpack/methods/ann/layer/noisylinear.hpp b/src/mlpack/methods/ann/layer/noisylinear.hpp
index 2260b27dbe..bc8c413af8 100644
--- a/src/mlpack/methods/ann/layer/noisylinear.hpp
+++ b/src/mlpack/methods/ann/layer/noisylinear.hpp
@@ -47,11 +47,24 @@ class NoisyLinear
   NoisyLinear(const size_t inSize,
          const size_t outSize);
 
+  //! Copy constructor.
+  NoisyLinear(const NoisyLinear&);
+
   /*
    * Reset the layer parameter.
    */
   void Reset();
 
+  /*
+   * Reset the noise parameters(epsilons).
+   */
+  void ResetNoise();
+
+  /*
+   * Reset the values of layer parameters (factorized gaussian noise).
+   */
+  void ResetParameters();
+
   /**
    * Ordinary feed forward pass of a neural network, evaluating the function
    * f(x) by propagating the activity forward through f.
@@ -141,9 +154,27 @@ class NoisyLinear
   //! Locally-stored weight parameters.
   OutputDataType weight;
 
-  //! Locally-stored bias term parameters.
+  //! Locally-stored weight-mean parameters.
+  OutputDataType weightMu;
+
+  //! Locally-stored weight-standard-deviation parameters.
+  OutputDataType weightSigma;
+
+  //! Locally-stored weight-epsilon parameters.
+  OutputDataType weightEpsilon;
+
+  //! Locally-stored bias parameters.
   OutputDataType bias;
 
+  //! Locally-stored bias-mean parameters.
+  OutputDataType biasMu;
+
+  //! Locally-stored bias-standard-deviation parameters.
+  OutputDataType biasSigma;
+
+  //! Locally-stored bias-epsilon parameters.
+  OutputDataType biasEpsilon;
+
   //! Locally-stored delta object.
   OutputDataType delta;
 
diff --git a/src/mlpack/methods/ann/layer/noisylinear_impl.hpp b/src/mlpack/methods/ann/layer/noisylinear_impl.hpp
index 849f29278f..8c9d4c969a 100644
--- a/src/mlpack/methods/ann/layer/noisylinear_impl.hpp
+++ b/src/mlpack/methods/ann/layer/noisylinear_impl.hpp
@@ -26,6 +26,16 @@ NoisyLinear<InputDataType, OutputDataType>::NoisyLinear() :
   // Nothing to do here.
 }
 
+template<typename InputDataType, typename OutputDataType>
+NoisyLinear<InputDataType, OutputDataType>::NoisyLinear(
+  const NoisyLinear& layer) :
+    inSize(layer.inSize),
+    outSize(layer.outSize),
+    weights(layer.weights)
+{
+  Reset();
+}
+
 template<typename InputDataType, typename OutputDataType>
 NoisyLinear<InputDataType, OutputDataType>::NoisyLinear(
     const size_t inSize,
@@ -33,15 +43,40 @@ NoisyLinear<InputDataType, OutputDataType>::NoisyLinear(
     inSize(inSize),
     outSize(outSize)
 {
-  weights.set_size(outSize * inSize + outSize, 1);
+  weights.set_size((outSize * inSize + outSize) * 2, 1);
+  weightEpsilon.set_size(outSize, inSize);
+  biasEpsilon.set_size(outSize, 1);
 }
 
 template<typename InputDataType, typename OutputDataType>
 void NoisyLinear<InputDataType, OutputDataType>::Reset()
 {
-  weight = arma::mat(weights.memptr(), outSize, inSize, false, false);
-  bias = arma::mat(weights.memptr() + weight.n_elem,
+  weightMu = arma::mat(weights.memptr(),
+      outSize, inSize, false, false);
+  biasMu = arma::mat(weights.memptr() + weightMu.n_elem,
       outSize, 1, false, false);
+  weightSigma = arma::mat(weights.memptr() + weightMu.n_elem + biasMu.n_elem,
+      outSize, inSize, false, false);
+  biasSigma = arma::mat(weights.memptr() + weightMu.n_elem * 2 + biasMu.n_elem,
+      outSize, 1, false, false);
+  this->ResetNoise();
+}
+
+template<typename InputDataType, typename OutputDataType>
+void NoisyLinear<InputDataType, OutputDataType>::ResetNoise()
+{
+  // TODO: Resets noise parameters.
+  weightEpsilon.randu();
+  biasEpsilon.randu();
+}
+
+template<typename InputDataType, typename OutputDataType>
+void NoisyLinear<InputDataType, OutputDataType>::ResetParameters()
+{
+  /*
+  * TODO: Reset network parameters according to factorized gaussion 
+  * initialization
+  */
 }
 
 template<typename InputDataType, typename OutputDataType>
@@ -49,6 +84,8 @@ template<typename eT>
 void NoisyLinear<InputDataType, OutputDataType>::Forward(
     const arma::Mat<eT>& input, arma::Mat<eT>& output)
 {
+  weight = weightMu + weightSigma % weightEpsilon;
+  bias = biasMu + biasSigma % biasEpsilon;
   output = weight * input;
   output.each_col() += bias;
 }
@@ -68,10 +105,20 @@ void NoisyLinear<InputDataType, OutputDataType>::Gradient(
     const arma::Mat<eT>& error,
     arma::Mat<eT>& gradient)
 {
-  gradient.submat(0, 0, weight.n_elem - 1, 0) = arma::vectorise(
-      error * input.t());
-  gradient.submat(weight.n_elem, 0, gradient.n_elem - 1, 0) =
-      arma::sum(error, 1);
+  // Locally stored to prevent multiplication twice.
+  arma::mat weightGrad = error * input.t();
+
+  // Gradients for mu values
+  gradient.rows(0, weight.n_elem - 1)
+      = arma::vectorise(weightGrad);
+  gradient.rows(weight.n_elem, weight.n_elem + bias.n_elem - 1)
+      = arma::sum(error, 1);
+
+  // Gradients for sigma values
+  gradient.rows(weight.n_elem + bias.n_elem, gradient.n_elem - bias.n_elem - 1)
+      = arma::vectorise(weightGrad % weightEpsilon);
+  gradient.rows(gradient.n_elem - bias.n_elem, gradient.n_elem - 1)
+      = arma::sum(error % biasEpsilon, 1);
 }
 
 template<typename InputDataType, typename OutputDataType>
@@ -85,7 +132,7 @@ void NoisyLinear<InputDataType, OutputDataType>::serialize(
   // This is inefficient, but we have to allocate this memory so that
   // WeightSetVisitor gets the right size.
   if (Archive::is_loading::value)
-    weights.set_size(outSize * inSize + outSize, 1);
+    weights.set_size((outSize * inSize + outSize) * 2, 1);
 }
 
 } // namespace ann