updates to gerstner syn, data-loader

Author: Alexander Ororbia

ngclearn/components/synapses/hebbian/gerstnerHebbianSynapse.py

@@ -23,6 +23,10 @@ class GerstnerHebbianSynapse(DenseSynapse):
     | c2_corr < 0 and c0 = c1_pre = c1_post = 0 => anti-Hebbian update
     | c2_corr = 1 and c1_pre = -x_theta < 0
 
+    | References: 
+    | Gerstner, W. and Kistler, W.M., 2002. Mathematical formulations of Hebbian 
+    | learning. Biological cybernetics, 87(5), pp.404-415.
+
     """
     def __init__(
         self,
@@ -37,7 +41,7 @@ def __init__(
         batch_size=1,
         **kwargs
     ):
-        bias_init = None ## no biases are included in Gerster's formulation
+        bias_init = None ## NOTE: no biases are included in Gerster's formulation
         super().__init__(
             name,
             shape=shape,
@@ -48,11 +52,11 @@ def __init__(
             batch_size=batch_size,
             **kwargs
         )
-        ## General Hebbian meta-parameters
+        ## general Hebbian meta-parameters
         self.eta = eta
         self.sign_value = sign_value
 
-        ## Expansion coefficients (c0, c1_pre, c1_post, c2_corr)
+        ## Gerstner and Kisler's expansion coefficients (c0, c1_pre, c1_post, c2_corr)
         if coeffs is None: ## Default to standard bilinear Hebb
             self.coeffs = {
                 'c0': 0., 'c1_pre': 0., 'c1_post': 0., 'c2_corr': 1.0
@@ -64,50 +68,43 @@ def __init__(
         self.c1_post = self.coeffs['c1_post']
         self.c2_corr = self.coeffs['c2_corr']
 
-        # Initialize Weights (using JAX PRNG)
-        #init_key, _ = random.split(self.key)
-        #w_init = random.normal(init_key, shape) * 0.05
-
-        # Compartments (ngc-learn state management)
-        #self.weights = Compartment(w_init)
+        ## set up relevant compartments
         self.pre = Compartment(jnp.zeros((1, shape[1])))
         self.post = Compartment(jnp.zeros((1, shape[0])))
+        self.dWeights = Compartment(jnp.zeros(shape))
 
     @compilable
-    def evolve(self, **kwargs):
-        """
-        Updates weights using the Gerstner general expansion.
-        Assumes pre_act and post_act compartments have been populated.
-        """
-        # Retrieve current states
+    def evolve(self, **kwargs): ## perform update via Gerstner's general expansion
+        ## retrieve current compartment state values
         W = self.weights.get()
-        x = self.pre.get()  # pre-synaptic activity (batch, pre_dim)
-        y = self.post.get() # post-synaptic activity (batch, post_dim)
+        x = self.pre.get()  ## pre-synaptic activity (batch, pre_dim)
+        y = self.post.get() ## post-synaptic activity (batch, post_dim)
         batch_size = self.batch_size
 
-        ## Bilinear Term (c2): correlation matrix
-        ### (post_dim, batch) @ (batch, pre_dim) -> (post_dim, pre_dim)
+        ## calculate bilinear Term (c2), i.e., correlation matrix
+        ### (pre_dim, batch) @ (batch, post_dim) -> (pre_dim, post_dim)
         dW_corr = jnp.matmul(x.T, y) * (1./batch_size)
-        ## Linear pre-synaptic term (c1_pre)
-        ### Average over batch then broadcast to match weight matrix
+        ## linear pre-synaptic term (c1_pre)
+        ### get mean over batch then broadcast to match weight matrix
         dW_pre = jnp.sum(x, axis=0, keepdims=True).T * (1./batch_size)
-        ## Linear post-synaptic term (c1_post)
-        dW_post = jnp.sum(y, axis=0, keepdims=True) * (1./batch_size)
+        ## linear post-synaptic term (c1_post), mean over post-syn values
+        dW_post = jnp.sum(y, axis=0, keepdims=True) * (1./batch_size) 
 
-        ## Apply Equation 3 Taylor expansion
+        ## apply Taylor expansion from Equation 3 (Gerstner and Kistler)
         dW = (self.c0 * W +  ## synaptic decay
               self.c1_pre * dW_pre +  ## bilinear term
               self.c1_post * dW_post +  ## pre-synaptic gating term
               self.c2_corr * dW_corr  ## post-synpatic gating term
         )
+        self.dWeights.set(dW)
+
         ## perform a step of Hebbian ascent
-        W = W + self.eta * dW
-        ## Update weights
+        W = W + self.eta * dW ## update synaptic efficacies
         self.weights.set(W)
 
     @compilable
-    def reset(self, **kwargs):
-        """Clears activity compartments"""
+    def reset(self, **kwargs): ## clear compartment values
         self.pre.set( jnp.zeros((self.batch_size, self.shape[1])) )
         self.post.set( jnp.zeros((self.batch_size, self.shape[0])) )
+        self.dWeights.set(self.dWeights.get() * 0)
 

ngclearn/utils/data_loader.py

@@ -25,8 +25,14 @@ class DataLoader(object):
         key: PRNG key to control determinism of any underlying random values
             associated with this synaptic cable
     """
-    def __init__(self, design_matrices, batch_size, disable_shuffle=False,
-                 ensure_equal_batches=True, key=None):
+    def __init__(
+            self,
+            design_matrices,
+            batch_size,
+            disable_shuffle=False,
+            ensure_equal_batches=True,
+            key=None
+    ):
         self.key = key
         if self.key is None:
             self.key = random.PRNGKey(time.time_ns())
@@ -47,23 +53,25 @@ def __init__(self, design_matrices, batch_size, disable_shuffle=False,
 
     def __iter__(self):
         """
-        Yields a mini-batch of the form:  [("name", batch),("name",batch),...]
+        Yields a mini-batch of the form:
+
+        | batch = [("name", batchx), ("name", batchy),...("name", batchz)]
         """
-        if self.disable_shuffle == False:
+        if not self.disable_shuffle: #self.disable_shuffle == False:
             self.key, *subkeys = random.split(self.key, 2)
             self.ptrs = random.permutation(subkeys[0], self.data_len)
         idx = 0
-        while idx < len(self.ptrs): # go through each sample via the sampling pointer
+        while idx < len(self.ptrs): ## go through each sample via the sampling pointer
             e_idx = idx + self.batch_size
-            if e_idx > len(self.ptrs): # prevents reaching beyond length of dataset
+            if e_idx > len(self.ptrs): ## prevents reaching beyond length of dataset
                 e_idx = len(self.ptrs)
-            # extract sampling integer pointers
+            ## extract sampling integer pointers
             indices = self.ptrs[idx:e_idx]
-            if self.ensure_equal_batches == True:
+            if self.ensure_equal_batches: # == True:
                 if indices.shape[0] < self.batch_size:
                     diff = self.batch_size - indices.shape[0]
                     indices = jnp.concatenate((indices, self.ptrs[0:diff]))
-            # create the actual pattern vector batch block matrices
+            ## create the actual pattern vector batch block matrices
             data_batch = []
             for dname, dmatrix in self.design_matrices:
                 x_batch = dmatrix[indices]