update to mstdp-et and var-trace

Author: Alexander Ororbia

ngclearn/components/other/varTrace.py

@@ -59,6 +59,8 @@ class VarTrace(JaxComponent): ## low-pass filter
         a_delta: value to increment a trace by in presence of a spike; note if set
             to a value <= 0, then a piecewise gated trace will be used instead
 
+        P_scale: if `a_delta=0`, then this scales the value that the trace snaps to upon receiving a pulse value
+
         gamma_tr: an extra multiplier in front of the leak of the trace (Default: 1)
 
         decay_type: string indicating the decay type to be applied to ODE
@@ -73,13 +75,14 @@ class VarTrace(JaxComponent): ## low-pass filter
     """
 
     # Define Functions
-    def __init__(self, name, n_units, tau_tr, a_delta, gamma_tr=1, decay_type="exp",
+    def __init__(self, name, n_units, tau_tr, a_delta, P_scale=1., gamma_tr=1, decay_type="exp",
                  batch_size=1, **kwargs):
         super().__init__(name, **kwargs)
 
         ## Trace control coefficients
         self.tau_tr = tau_tr ## trace time constant
         self.a_delta = a_delta ## trace increment (if spike occurred)
+        self.P_scale = P_scale ## trace scale if non-additive trace to be used
         self.gamma_tr = gamma_tr
         self.decay_type = decay_type ## lin --> linear decay; exp --> exponential decay
 
@@ -94,7 +97,7 @@ def __init__(self, name, n_units, tau_tr, a_delta, gamma_tr=1, decay_type="exp",
 
     @transition(output_compartments=["outputs", "trace"])
     @staticmethod
-    def advance_state(dt, decay_type, tau_tr, a_delta, gamma_tr, inputs, trace):
+    def advance_state(dt, decay_type, tau_tr, a_delta, P_scale, gamma_tr, inputs, trace):
         decayFactor = 0.
         if "exp" in decay_type:
             decayFactor = jnp.exp(-dt/tau_tr)
@@ -104,7 +107,7 @@ def advance_state(dt, decay_type, tau_tr, a_delta, gamma_tr, inputs, trace):
         if a_delta > 0.:
             _x_tr = _x_tr + inputs * a_delta
         else:
-            _x_tr = _x_tr * (1. - inputs) + inputs
+            _x_tr = _x_tr * (1. - inputs) + inputs * P_scale
         trace = _x_tr
         return trace, trace
 

ngclearn/components/synapses/modulated/MSTDPETSynapse.py

@@ -61,6 +61,8 @@ class MSTDPETSynapse(TraceSTDPSynapse): # modulated trace-based STDP w/ eligilit
 
         elg_decay: eligibility decay constant (default: 1)
 
+        tau_w: amount of synaptic decay to augment each MSTDP/MSTDP-ET update with
+
         weight_init: a kernel to drive initialization of this synaptic cable's values;
             typically a tuple with 1st element as a string calling the name of
             initialization to use
@@ -74,26 +76,28 @@ class MSTDPETSynapse(TraceSTDPSynapse): # modulated trace-based STDP w/ eligilit
 
     # Define Functions
     def __init__(
-            self, name, shape, A_plus, A_minus, eta=1., mu=0., pretrace_target=0., tau_elg=0., elg_decay=1.,
+            self, name, shape, A_plus, A_minus, eta=1., mu=0., pretrace_target=0., tau_elg=0., elg_decay=1., tau_w=0.,
             weight_init=None, resist_scale=1., p_conn=1., w_bound=1., batch_size=1, **kwargs
     ):
         super().__init__(
             name, shape, A_plus, A_minus, eta=eta, mu=mu, pretrace_target=pretrace_target, weight_init=weight_init,
             resist_scale=resist_scale, p_conn=p_conn, w_bound=w_bound, batch_size=batch_size, **kwargs
         )
         self.w_eps = 0.
+        self.tau_w = tau_w
         ## MSTDP/MSTDP-ET meta-parameters
         self.tau_elg = tau_elg
         self.elg_decay = elg_decay
         ## MSTDP/MSTDP-ET compartments
         self.modulator = Compartment(jnp.zeros((self.batch_size, 1)))
         self.eligibility = Compartment(jnp.zeros(shape))
+        self.outmask = Compartment(jnp.zeros((1, shape[1])))
 
     @transition(output_compartments=["weights", "dWeights", "eligibility"])
     @staticmethod
     def evolve(
-            dt, w_bound, w_eps, preTrace_target, mu, Aplus, Aminus, tau_elg, elg_decay, preSpike, postSpike, preTrace,
-            postTrace, weights, dWeights, eta, modulator, eligibility
+            dt, w_bound, w_eps, preTrace_target, mu, Aplus, Aminus, tau_elg, elg_decay, tau_w, preSpike, postSpike,
+            preTrace, postTrace, weights, dWeights, eta, modulator, eligibility, outmask
     ):
         # dW_dt = TraceSTDPSynapse._compute_update( ## use Hebbian/STDP rule to obtain a non-modulated update
         #     dt, w_bound, preTrace_target, mu, Aplus, Aminus, preSpike, postSpike, preTrace, postTrace, weights
@@ -105,21 +109,25 @@ def evolve(
         else: ## otherwise, just do M-STDP
             eligibility = dWeights ## dynamics of M-STDP had no eligibility tracing
         ## do a gradient ascent update/shift
-        weights = weights + eligibility * modulator * eta  ## do modulated update
-        #'''
+        decayTerm = 0.
+        if tau_w > 0.:
+            decayTerm = weights * (1. / tau_w)
+        weights = weights + (eligibility * modulator * eta) * outmask - decayTerm ## do modulated update
+
         dW_dt = TraceSTDPSynapse._compute_update( ## use Hebbian/STDP rule to obtain a non-modulated update
             dt, w_bound, preTrace_target, mu, Aplus, Aminus, preSpike, postSpike, preTrace, postTrace, weights
         )
         dWeights = dW_dt ## can think of this as eligibility at time t
-        #'''
-    
+
         #w_eps = 0.01
         weights = jnp.clip(weights, w_eps, w_bound - w_eps)  # jnp.abs(w_bound))
 
         return weights, dWeights, eligibility
 
     @transition(
-        output_compartments=["inputs", "outputs", "preSpike", "postSpike", "preTrace", "postTrace", "dWeights", "eligibility"]
+        output_compartments=[
+            "inputs", "outputs", "preSpike", "postSpike", "preTrace", "postTrace", "dWeights", "eligibility", "outmask"
+        ]
     )
     @staticmethod
     def reset(batch_size, shape):
@@ -134,7 +142,8 @@ def reset(batch_size, shape):
         postTrace = postVals
         dWeights = synVals
         eligibility = synVals
-        return inputs, outputs, preSpike, postSpike, preTrace, postTrace, dWeights, eligibility
+        outmask = postVals + 1.
+        return inputs, outputs, preSpike, postSpike, preTrace, postTrace, dWeights, eligibility, outmask
 
     @classmethod
     def help(cls): ## component help function