weight_optimizer.h

Go to the documentation of this file.

/*
 *  weight_optimizer.h
 *
 *  This file is part of NEST.
 *
 *  Copyright (C) 2004 The NEST Initiative
 *
 *  NEST is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation, either version 2 of the License, or
 *  (at your option) any later version.
 *
 *  NEST is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with NEST.  If not, see <http://www.gnu.org/licenses/>.
 *
 */
 
#ifndef WEIGHT_OPTIMIZER_H
#define WEIGHT_OPTIMIZER_H
 
// nestkernel
#include "dictionary.h"
 
namespace nest
{
 
/* BeginUserDocs: e-prop plasticity, synapse
 
Short description
+++++++++++++++++
 
Selection of weight optimizers
 
Description
+++++++++++
A weight optimizer is an algorithm that adjusts the synaptic weights in a
network during training to minimize the loss function and thus improve the
network's performance on a given task.
 
This method is an essential part of plasticity rules like e-prop plasticity.
 
Currently two weight optimizers are implemented: gradient descent and the Adam optimizer.
 
In gradient descent :footcite:p:`Huh2018` the weights are optimized via:
 
.. math::
  W_t = W_{t-1} - \eta g_t \,, \\
 
where :math:`\eta` denotes the learning rate and :math:`g_t` the gradient of the current
time step :math:`t`.
 
In the Adam scheme :footcite:p:`Kingma2015` the weights are optimized via:
 
.. math::
  m_0 &= 0, v_0 = 0, t = 1 \,, \\
  m_t &= \beta_1 m_{t-1} + \left( 1- \beta_1 \right) g_t \,, \\
  v_t &= \beta_2 v_{t-1} + \left( 1 - \beta_2 \right) g_t^2 \,, \\
  \alpha_t &= \eta \frac{ \sqrt{ 1- \beta_2^t } }{ 1 - \beta_1^t } \,, \\
  W_t &= W_{t-1} - \alpha_t \frac{ m_t }{ \sqrt{v_t} + \hat{\epsilon} } \,. \\
 
Note that the implementation follows the implementation in TensorFlow :footcite:p:`KerasTeam2024` for comparability.
The TensorFlow implementation deviates from :footcite:p:`Huh2018` in that it assumes
:math:`\hat{\epsilon} = \epsilon \sqrt{ 1 - \beta_2^t }` to be constant, whereas :footcite:p:`Huh2018`
assumes :math:`\epsilon = \hat{\epsilon} \sqrt{ 1 - \beta_2^t }` to be constant.
 
When `optimize_each_step` is set to `True`, the weights are optimized at every
time step. If set to `False`, optimization occurs once per spike, resulting in a
significant speed-up. For gradient descent, both settings yield the same
results under exact arithmetic; however, small numerical differences may be
observed due to floating point precision. For the Adam optimizer, only setting
`optimize_each_step` to `True` precisely implements the algorithm as described
in :footcite:p:`Kingma2015`. The impact of this setting on learning performance may vary depending
on the task.
 
Parameters
++++++++++
 
The following parameters can be set in the status dictionary.
 
====================== ==== ========================= ========= ============================================
**Common optimizer parameters**
------------------------------------------------------------------------------------------------------------
Parameter              Unit Math equivalent           Default   Description
====================== ==== ========================= ========= ============================================
``batch_size``                                              1   Size of batch
``eta``                     :math:`\eta`                 1e-4   Learning rate
``optimize_each_step``                                 ``True`` If ``True``, optimize each step, if ``False``
                                                                once per spike
``Wmax``                pA  :math:`W_{ji}^\text{max}`   100.0   Maximal value for synaptic weight
``Wmin``                pA  :math:`W_{ji}^\text{min}`  -100.0   Minimal value for synaptic weight
====================== ==== ========================= ========= ============================================
 
========= ==== =============== ================== ==============
**Gradient descent parameters (default optimizer)**
----------------------------------------------------------------
Parameter Unit Math equivalent Default            Description
========= ==== =============== ================== ==============
``type``                       "gradient_descent" Optimizer type
========= ==== =============== ================== ==============
 
=========== ==== ================ ======= =================================================
**Adam optimizer parameters**
-------------------------------------------------------------------------------------------
Parameter   Unit Math equivalent  Default Description
=========== ==== ================ ======= =================================================
``type``                           "adam" Optimizer type
``beta_1``       :math:`\beta_1`      0.9 Exponential decay rate for first moment estimate
``beta_2``       :math:`\beta_2`    0.999 Exponential decay rate for second moment estimate
``epsilon``      :math:`\epsilon`    1e-7 Small constant for numerical stability
=========== ==== ================ ======= =================================================
 
The following state variables evolve during simulation.
 
============== ==== =============== ============= ==========================
**Adam optimizer state variables for individual synapses**
----------------------------------------------------------------------------
State variable Unit Math equivalent Initial value Description
============== ==== =============== ============= ==========================
``m``               :math:`m`                 0.0 First moment estimate
``v``               :math:`v`                 0.0 Second moment raw estimate
============== ==== =============== ============= ==========================
 
 
References
++++++++++
 
.. footbibliography::
 
See also
++++++++
 
Examples using this model
+++++++++++++++++++++++++
 
.. listexamples:: eprop_synapse_bsshslm_2020
 
EndUserDocs */
 
class WeightOptimizer;
 
class WeightOptimizerCommonProperties
{
public:
  WeightOptimizerCommonProperties();
 
  virtual ~WeightOptimizerCommonProperties()
  {
  }
 
  WeightOptimizerCommonProperties( const WeightOptimizerCommonProperties& );
 
  WeightOptimizer& operator=( const WeightOptimizer& ) = delete;
 
  virtual void get_status( Dictionary& d ) const;
 
  virtual void set_status( const Dictionary& d );
 
  virtual WeightOptimizerCommonProperties* clone() const = 0;
 
  virtual WeightOptimizer* get_optimizer() const = 0;
 
  double
  get_Wmin() const
  {
    return Wmin_;
  }
 
  double
  get_Wmax() const
  {
    return Wmax_;
  }
 
  virtual std::string get_name() const = 0;
 
public:
  size_t batch_size_;
 
  double eta_;
 
  double eta_first_change_;
 
  long eta_change_count_;
 
  double Wmin_;
 
  double Wmax_;
 
  bool optimize_each_step_;
};
 
class WeightOptimizer
{
public:
  WeightOptimizer();
 
  virtual ~WeightOptimizer()
  {
  }
 
  WeightOptimizer( const WeightOptimizer& ) = default;
 
  WeightOptimizer& operator=( const WeightOptimizer& ) = delete;
 
  virtual void get_status( Dictionary& d ) const;
 
  virtual void set_status( const Dictionary& d );
 
  double optimized_weight( const WeightOptimizerCommonProperties& cp,
    const size_t idx_current_update,
    const double gradient,
    double weight );
 
protected:
  virtual double optimize_( const WeightOptimizerCommonProperties& cp, double weight, size_t current_opt_step ) = 0;
 
  double cumulative_gradient_;
 
  size_t optimization_step_;
 
  double eta_current_;
 
  long n_optimize_;
};
 
class WeightOptimizerGradientDescent : public WeightOptimizer
{
public:
  WeightOptimizerGradientDescent();
 
  WeightOptimizerGradientDescent( const WeightOptimizerGradientDescent& ) = default;
 
  WeightOptimizerGradientDescent& operator=( const WeightOptimizerGradientDescent& ) = delete;
 
private:
  double optimize_( const WeightOptimizerCommonProperties& cp, double weight, size_t current_opt_step ) override;
};
 
class WeightOptimizerCommonPropertiesGradientDescent : public WeightOptimizerCommonProperties
{
  friend class WeightOptimizerGradientDescent;
 
public:
  WeightOptimizerCommonPropertiesGradientDescent();
 
  WeightOptimizerCommonPropertiesGradientDescent( const WeightOptimizerCommonPropertiesGradientDescent& ) = default;
 
  WeightOptimizerCommonPropertiesGradientDescent& operator=(
    const WeightOptimizerCommonPropertiesGradientDescent& ) = delete;
 
  WeightOptimizerCommonProperties* clone() const override;
  WeightOptimizer* get_optimizer() const override;
 
  std::string
  get_name() const override
  {
    return "gradient_descent";
  }
};
 
class WeightOptimizerAdam : public WeightOptimizer
{
public:
  WeightOptimizerAdam();
 
  WeightOptimizerAdam( const WeightOptimizerAdam& ) = default;
 
  WeightOptimizerAdam& operator=( const WeightOptimizerAdam& ) = delete;
 
  void get_status( Dictionary& d ) const override;
  void set_status( const Dictionary& d ) override;
 
private:
  double optimize_( const WeightOptimizerCommonProperties& cp, double weight, size_t current_opt_step ) override;
 
  double m_;
 
  double v_;
 
  double beta_1_power_;
 
  double beta_2_power_;
};
 
class WeightOptimizerCommonPropertiesAdam : public WeightOptimizerCommonProperties
{
  friend class WeightOptimizerAdam;
 
public:
  WeightOptimizerCommonPropertiesAdam();
 
  WeightOptimizerCommonPropertiesAdam( const WeightOptimizerCommonPropertiesAdam& ) = default;
 
  WeightOptimizerCommonPropertiesAdam& operator=( const WeightOptimizerCommonPropertiesAdam& ) = delete;
 
  WeightOptimizerCommonProperties* clone() const override;
  WeightOptimizer* get_optimizer() const override;
 
  void get_status( Dictionary& d ) const override;
  void set_status( const Dictionary& d ) override;
 
  std::string
  get_name() const override
  {
    return "adam";
  }
 
private:
  double beta_1_;
 
  double beta_2_;
 
  double epsilon_;
};
 
}  // namespace nest
 
#endif  // WEIGHT_OPTIMIZER_H