From 05ea912eb14ec0e29c5ad9f1fb2b57815df5c384 Mon Sep 17 00:00:00 2001 From: ago109 Date: Wed, 24 Jul 2024 13:10:54 -0400 Subject: [PATCH 01/61] modded bernoulli-cell to include max-frequency constraint --- .../input_encoders/bernoulliCell.py | 47 ++++++++++++++++--- 1 file changed, 41 insertions(+), 6 deletions(-) diff --git a/ngclearn/components/input_encoders/bernoulliCell.py b/ngclearn/components/input_encoders/bernoulliCell.py index da090bfb..61e74031 100755 --- a/ngclearn/components/input_encoders/bernoulliCell.py +++ b/ngclearn/components/input_encoders/bernoulliCell.py @@ -2,6 +2,7 @@ from ngclearn.components.jaxComponent import JaxComponent from jax import numpy as jnp, random, jit from ngclearn.utils import tensorstats +from functools import partial @jit def _update_times(t, s, tols): @@ -37,9 +38,33 @@ def _sample_bernoulli(dkey, data): s_t = random.bernoulli(dkey, p=data).astype(jnp.float32) return s_t +@partial(jit, static_argnums=[3]) +def _sample_constrained_bernoulli(dkey, data, dt, fmax=63.75): + """ + Samples a Bernoulli spike train on-the-fly that is constrained to emit + at a particular rate over a time window. + + Args: + dkey: JAX key to drive stochasticity/noise + + data: sensory data (vector/matrix) + + dt: integration time constant + + fmax: maximum frequency (Hz) + + Returns: + binary spikes + """ + pspike = data * (dt/1000.) * fmax + eps = random.uniform(dkey, data.shape, minval=0., maxval=1., dtype=jnp.float32) + s_t = (eps < pspike).astype(jnp.float32) + return s_t + class BernoulliCell(JaxComponent): """ - A Bernoulli cell that produces Bernoulli-distributed spikes on-the-fly. + A Bernoulli cell that produces variations of Bernoulli-distributed spikes + on-the-fly (including constrained-rate trains). | --- Cell Input Compartments: --- | inputs - input (takes in external signals) @@ -53,12 +78,17 @@ class BernoulliCell(JaxComponent): name: the string name of this cell n_units: number of cellular entities (neural population size) + + max_freq: maximum frequency (in Hertz) of this Bernoulli spike train (must be > 0.) """ # Define Functions - def __init__(self, name, n_units, batch_size=1, **kwargs): + def __init__(self, name, n_units, max_freq=63.75, batch_size=1, **kwargs): super().__init__(name, **kwargs) + ## Constrained Bernoulli meta-parameters + self.max_freq = max_freq ## maximum frequency (in Hertz/Hz) + ## Layer Size Setup self.batch_size = batch_size self.n_units = n_units @@ -70,11 +100,16 @@ def __init__(self, name, n_units, batch_size=1, **kwargs): self.tols = Compartment(restVals, display_name="Time-of-Last-Spike", units="ms") # time of last spike @staticmethod - def _advance_state(t, key, inputs, tols): + def _advance_state(t, dt, max_freq, key, inputs, tols): key, *subkeys = random.split(key, 2) - outputs = _sample_bernoulli(subkeys[0], data=inputs) - timeOfLastSpike = _update_times(t, outputs, tols) - return outputs, timeOfLastSpike, key + if max_freq > 0.: + outputs = _sample_constrained_bernoulli( ## sample Bernoulli w/ target rate + subkeys[0], data=inputs, dt=dt, fmax=max_freq + ) + else: + outputs = _sample_bernoulli(subkeys[0], data=inputs) + tols = _update_times(t, outputs, tols) + return outputs, tols, key @resolver(_advance_state) def advance_state(self, outputs, tols, key): From c19d15ebf5f0922e9e807652845cda8685c92911 Mon Sep 17 00:00:00 2001 From: ago109 Date: Wed, 24 Jul 2024 14:41:36 -0400 Subject: [PATCH 02/61] added warning check to bernoulli, some cleanup --- .../input_encoders/bernoulliCell.py | 30 ++++++++++++++----- 1 file changed, 23 insertions(+), 7 deletions(-) diff --git a/ngclearn/components/input_encoders/bernoulliCell.py b/ngclearn/components/input_encoders/bernoulliCell.py index 61e74031..fa036302 100755 --- a/ngclearn/components/input_encoders/bernoulliCell.py +++ b/ngclearn/components/input_encoders/bernoulliCell.py @@ -3,6 +3,8 @@ from jax import numpy as jnp, random, jit from ngclearn.utils import tensorstats from functools import partial +from ngcsimlib.deprecators import deprecate_args +from ngcsimlib.logger import info, warn @jit def _update_times(t, s, tols): @@ -79,15 +81,15 @@ class BernoulliCell(JaxComponent): n_units: number of cellular entities (neural population size) - max_freq: maximum frequency (in Hertz) of this Bernoulli spike train (must be > 0.) + target_freq: maximum frequency (in Hertz) of this Bernoulli spike train (must be > 0.) """ - # Define Functions - def __init__(self, name, n_units, max_freq=63.75, batch_size=1, **kwargs): + @deprecate_args(target_freq="max_freq") + def __init__(self, name, n_units, target_freq=63.75, batch_size=1, **kwargs): super().__init__(name, **kwargs) ## Constrained Bernoulli meta-parameters - self.max_freq = max_freq ## maximum frequency (in Hertz/Hz) + self.target_freq = target_freq ## maximum frequency (in Hertz/Hz) ## Layer Size Setup self.batch_size = batch_size @@ -99,12 +101,26 @@ def __init__(self, name, n_units, max_freq=63.75, batch_size=1, **kwargs): self.outputs = Compartment(restVals, display_name="Spikes") # output compartment self.tols = Compartment(restVals, display_name="Time-of-Last-Spike", units="ms") # time of last spike + def validate(self, dt, **validation_kwargs): + ## check for unstable combinations of dt and target-frequency meta-params + valid = super().validate(**validation_kwargs) + events_per_timestep = (dt/1000.) * self.target_freq ## compute scaled probability + if events_per_timestep > 1.: + valid = False + warn( + f"{self.name} will be unable to make as many temporal events as " + f"requested! ({events_per_timestep} events/timestep) Unstable " + f"combination of dt = {dt} and target_freq = {self.target_freq} " + f"being used!" + ) + return valid + @staticmethod - def _advance_state(t, dt, max_freq, key, inputs, tols): + def _advance_state(t, dt, target_freq, key, inputs, tols): key, *subkeys = random.split(key, 2) - if max_freq > 0.: + if target_freq > 0.: outputs = _sample_constrained_bernoulli( ## sample Bernoulli w/ target rate - subkeys[0], data=inputs, dt=dt, fmax=max_freq + subkeys[0], data=inputs, dt=dt, fmax=target_freq ) else: outputs = _sample_bernoulli(subkeys[0], data=inputs) From 23a54f60fe9a39910e7f65e763debb7c72f60ffb Mon Sep 17 00:00:00 2001 From: ago109 Date: Wed, 24 Jul 2024 16:00:28 -0400 Subject: [PATCH 03/61] integrated if-cell, cleaned up lif and inits --- ngclearn/components/__init__.py | 1 + ngclearn/components/jaxComponent.py | 1 - ngclearn/components/neurons/__init__.py | 1 + ngclearn/components/neurons/spiking/IFCell.py | 306 ++++++++++++++++++ .../components/neurons/spiking/LIFCell.py | 47 ++- .../components/neurons/spiking/__init__.py | 1 + 6 files changed, 345 insertions(+), 12 deletions(-) create mode 100755 ngclearn/components/neurons/spiking/IFCell.py diff --git a/ngclearn/components/__init__.py b/ngclearn/components/__init__.py index d9534871..005fbacc 100644 --- a/ngclearn/components/__init__.py +++ b/ngclearn/components/__init__.py @@ -6,6 +6,7 @@ from .neurons.graded.rewardErrorCell import RewardErrorCell ## point to standard spiking cell component types from .neurons.spiking.sLIFCell import SLIFCell +from .neurons.spiking.IFCell import IFCell from .neurons.spiking.LIFCell import LIFCell from .neurons.spiking.WTASCell import WTASCell from .neurons.spiking.quadLIFCell import QuadLIFCell diff --git a/ngclearn/components/jaxComponent.py b/ngclearn/components/jaxComponent.py index 8286c6c0..f07309fe 100755 --- a/ngclearn/components/jaxComponent.py +++ b/ngclearn/components/jaxComponent.py @@ -21,4 +21,3 @@ def __init__(self, name, key=None, directory=None, **kwargs): self.directory = directory self.key = Compartment( random.PRNGKey(time.time_ns()) if key is None else key) - diff --git a/ngclearn/components/neurons/__init__.py b/ngclearn/components/neurons/__init__.py index 900a58ce..42a4a971 100644 --- a/ngclearn/components/neurons/__init__.py +++ b/ngclearn/components/neurons/__init__.py @@ -5,6 +5,7 @@ from .graded.rewardErrorCell import RewardErrorCell ## point to standard spiking cell component types from .spiking.sLIFCell import SLIFCell +from .spiking.IFCell import IFCell from .spiking.LIFCell import LIFCell from .spiking.WTASCell import WTASCell from .spiking.quadLIFCell import QuadLIFCell diff --git a/ngclearn/components/neurons/spiking/IFCell.py b/ngclearn/components/neurons/spiking/IFCell.py new file mode 100755 index 00000000..2c9acf52 --- /dev/null +++ b/ngclearn/components/neurons/spiking/IFCell.py @@ -0,0 +1,306 @@ +from jax import numpy as jnp, random, jit, nn +from ngclearn.utils import tensorstats +from ngcsimlib.deprecators import deprecate_args +from ngclearn import resolver, Component, Compartment +from ngclearn.components.jaxComponent import JaxComponent +from ngclearn.utils.diffeq.ode_utils import get_integrator_code, \ + step_euler, step_rk2 +from ngclearn.utils.surrogate_fx import (arctan_estimator, + triangular_estimator, + straight_through_estimator) + +@jit +def _update_times(t, s, tols): + """ + Updates time-of-last-spike (tols) variable. + + Args: + t: current time (a scalar/int value) + + s: binary spike vector + + tols: current time-of-last-spike variable + + Returns: + updated tols variable + """ + _tols = (1. - s) * tols + (s * t) + return _tols + +@jit +def _dfv_internal(j, v, rfr, tau_m, refract_T): ## raw voltage dynamics + mask = (rfr >= refract_T).astype(jnp.float32) # get refractory mask + ## update voltage / membrane potential + dv_dt = (j * mask) ## integration only involves electrical current + dv_dt = dv_dt * (1./tau_m) + return dv_dt + +def _dfv(t, v, params): ## voltage dynamics wrapper + j, rfr, tau_m, refract_T = params + dv_dt = _dfv_internal(j, v, rfr, tau_m, refract_T) + return dv_dt + +def _run_cell(dt, j, v, v_thr, rfr, tau_m, v_rest, v_reset, refract_T, integType=0): + ### Runs integrator (or integrate-and-fire; IF) neuronal dynamics + ## update voltage / membrane potential + v_params = (j, rfr, tau_m, refract_T) + if integType == 1: + _, _v = step_rk2(0., v, _dfv, dt, v_params) + else: + _, _v = step_euler(0., v, _dfv, dt, v_params) + ## obtain action potentials/spikes + s = (_v > v_thr).astype(jnp.float32) + ## update refractory variables + _rfr = (rfr + dt) * (1. - s) + ## perform hyper-polarization of neuronal cells + _v = _v * (1. - s) + s * v_reset + return _v, s, _rfr + +class IFCell(JaxComponent): ## integrate-and-fire cell + """ + A spiking cell based on integrate-and-fire (IF) neuronal dynamics. + + The specific differential equation that characterizes this cell + is (for adjusting v, given current j, over time) is: + + | tau_m * dv/dt = (v_rest - v) + j * R + | where R is the membrane resistance and v_rest is the resting potential + | also, if a spike occurs, v is set to v_reset + + | --- Cell Input Compartments: --- + | j - electrical current input (takes in external signals) + | --- Cell State Compartments: --- + | v - membrane potential/voltage state + | rfr - (relative) refractory variable state + | key - JAX PRNG key + | --- Cell Output Compartments: --- + | s - emitted binary spikes/action potentials + | s_raw - raw spike signals before post-processing (only if one_spike = True, else s_raw = s) + | tols - time-of-last-spike + + Args: + name: the string name of this cell + + n_units: number of cellular entities (neural population size) + + tau_m: membrane time constant + + resist_m: membrane resistance value (default: 1) + + thr: base value for adaptive thresholds that govern short-term + plasticity (in milliVolts, or mV; default: -52. mV) + + v_rest: membrane resting potential (in mV; default: -65 mV) + + v_reset: membrane reset potential (in mV) -- upon occurrence of a spike, + a neuronal cell's membrane potential will be set to this value; + (default: -60 mV) + + refract_time: relative refractory period time (ms; default: 0 ms) + + integration_type: type of integration to use for this cell's dynamics; + current supported forms include "euler" (Euler/RK-1 integration) + and "midpoint" or "rk2" (midpoint method/RK-2 integration) (Default: "euler") + + :Note: setting the integration type to the midpoint method will + increase the accuray of the estimate of the cell's evolution + at an increase in computational cost (and simulation time) + + surrgoate_type: type of surrogate function to use for approximating a + partial derivative of this cell's spikes w.r.t. its voltage/current + (default: "straight_through") + + :Note: surrogate options available include: "straight_through" + (straight-through estimator), "triangular" (triangular estimator), + and "arctan" (arc-tangent estimator) + + lower_clamp_voltage: if True, this will ensure voltage never is below + the value of `v_rest` (default: True) + """ + + @deprecate_args(thr_jitter=None) + def __init__(self, name, n_units, tau_m, resist_m=1., thr=-52., v_rest=-65., + v_reset=-60., refract_time=0., integration_type="euler", + surrgoate_type="straight_through", lower_clamp_voltage=True, + **kwargs): + super().__init__(name, **kwargs) + + ## Integration properties + self.integrationType = integration_type + self.intgFlag = get_integrator_code(self.integrationType) + + ## membrane parameter setup (affects ODE integration) + self.tau_m = tau_m ## membrane time constant + self.resist_m = resist_m ## resistance value + + self.v_rest = v_rest #-65. # mV + self.v_reset = v_reset # -60. # -65. # mV (milli-volts) + ## basic asserts to prevent neuronal dynamics breaking... + assert self.resist_m > 0. + self.refract_T = refract_time #5. # 2. ## refractory period # ms + self.thr = thr ## (fixed) base value for threshold #-52 # -72. # mV + self.lower_clamp_voltage = lower_clamp_voltage + + ## Layer Size Setup + self.batch_size = 1 + self.n_units = n_units + + ## set up surrogate function for spike emission + if surrgoate_type == "arctan": + self.spike_fx, self.d_spike_fx = arctan_estimator() + elif surrgoate_type == "triangular": + self.spike_fx, self.d_spike_fx = triangular_estimator() + else: ## default: straight_through + self.spike_fx, self.d_spike_fx = straight_through_estimator() + + + ## Compartment setup + restVals = jnp.zeros((self.batch_size, self.n_units)) + self.j = Compartment(restVals, display_name="Current", units="mA") + self.v = Compartment(restVals + self.v_rest, + display_name="Voltage", units="mV") + self.s = Compartment(restVals, display_name="Spikes") + self.rfr = Compartment(restVals + self.refract_T, + display_name="Refractory Time Period", units="ms") + self.tols = Compartment(restVals, display_name="Time-of-Last-Spike", + units="ms") ## time-of-last-spike + self.surrogate = Compartment(restVals + 1., display_name="Surrogate State Value") + + @staticmethod + def _advance_state(t, dt, tau_m, resist_m, v_rest, v_reset, refract_T, + thr, lower_clamp_voltage, intgFlag, d_spike_fx, key, + j, v, rfr, tols): + ## run one integration step for neuronal dynamics + j = j * resist_m + v, s, rfr = _run_cell(dt, j, v, thr, rfr, tau_m, v_rest, v_reset, + refract_T, intgFlag) + surrogate = d_spike_fx(v, thr) + ## update tols + tols = _update_times(t, s, tols) + if lower_clamp_voltage: ## ensure voltage never < v_rest + v = jnp.maximum(v, v_rest) + return v, s, rfr, tols, key, surrogate + + @resolver(_advance_state) + def advance_state(self, v, s, rfr, tols, key, surrogate): + self.v.set(v) + self.s.set(s) + self.rfr.set(rfr) + self.tols.set(tols) + self.key.set(key) + self.surrogate.set(surrogate) + + @staticmethod + def _reset(batch_size, n_units, v_rest, refract_T): + restVals = jnp.zeros((batch_size, n_units)) + j = restVals #+ 0 + v = restVals + v_rest + s = restVals #+ 0 + rfr = restVals + refract_T + tols = restVals #+ 0 + surrogate = restVals + 1. + return j, v, s, rfr, tols, surrogate + + @resolver(_reset) + def reset(self, j, v, s, rfr, tols, surrogate): + self.j.set(j) + self.v.set(v) + self.s.set(s) + self.rfr.set(rfr) + self.tols.set(tols) + self.surrogate.set(surrogate) + + def save(self, directory, **kwargs): + ## do a protected save of constants, depending on whether they are floats or arrays + tau_m = (self.tau_m if isinstance(self.tau_m, float) + else jnp.ones([[self.tau_m]])) + thr = (self.thr if isinstance(self.thr, float) + else jnp.ones([[self.thr]])) + v_rest = (self.v_rest if isinstance(self.v_rest, float) + else jnp.ones([[self.v_rest]])) + v_reset = (self.v_reset if isinstance(self.v_reset, float) + else jnp.ones([[self.v_reset]])) + v_decay = (self.v_decay if isinstance(self.v_decay, float) + else jnp.ones([[self.v_decay]])) + resist_m = (self.resist_m if isinstance(self.resist_m, float) + else jnp.ones([[self.resist_m]])) + tau_theta = (self.tau_theta if isinstance(self.tau_theta, float) + else jnp.ones([[self.tau_theta]])) + theta_plus = (self.theta_plus if isinstance(self.theta_plus, float) + else jnp.ones([[self.theta_plus]])) + + file_name = directory + "/" + self.name + ".npz" + jnp.savez(file_name, + tau_m=tau_m, thr=thr, v_rest=v_rest, + v_reset=v_reset, v_decay=v_decay, + resist_m=resist_m, tau_theta=tau_theta, + theta_plus=theta_plus, + key=self.key.value) + + def load(self, directory, seeded=False, **kwargs): + file_name = directory + "/" + self.name + ".npz" + data = jnp.load(file_name) + ## constants loaded in + self.tau_m = data['tau_m'] + self.thr = data['thr'] + self.v_rest = data['v_rest'] + self.v_reset = data['v_reset'] + self.v_decay = data['v_decay'] + self.resist_m = data['resist_m'] + self.tau_theta = data['tau_theta'] + self.theta_plus = data['theta_plus'] + + if seeded: + self.key.set(data['key']) + + @classmethod + def help(cls): ## component help function + properties = { + "cell_type": "IFCell - evolves neurons according to integrate-" + "and-fire spiking dynamics." + } + compartment_props = { + "inputs": + {"j": "External input electrical current"}, + "states": + {"v": "Membrane potential/voltage at time t", + "rfr": "Current state of (relative) refractory variable", + "thr": "Current state of voltage threshold at time t", + "key": "JAX PRNG key"}, + "outputs": + {"s": "Emitted spikes/pulses at time t", + "tols": "Time-of-last-spike"}, + } + hyperparams = { + "n_units": "Number of neuronal cells to model in this layer", + "tau_m": "Cell membrane time constant", + "resist_m": "Membrane resistance value", + "thr": "Base voltage threshold value", + "v_rest": "Resting membrane potential value", + "v_reset": "Reset membrane potential value", + "refract_time": "Length of relative refractory period (ms)", + "integration_type": "Type of numerical integration to use for the cell dynamics", + "surrgoate_type": "Type of surrogate function to use approximate " + "derivative of spike w.r.t. voltage/current", + "lower_bound_clamp": "Should voltage be lower bounded to be never be below `v_rest`" + } + info = {cls.__name__: properties, + "compartments": compartment_props, + "dynamics": "tau_m * dv/dt = (v_rest - v) + j * resist_m", + "hyperparameters": hyperparams} + return info + + def __repr__(self): + comps = [varname for varname in dir(self) if Compartment.is_compartment(getattr(self, varname))] + maxlen = max(len(c) for c in comps) + 5 + lines = f"[{self.__class__.__name__}] PATH: {self.name}\n" + for c in comps: + stats = tensorstats(getattr(self, c).value) + if stats is not None: + line = [f"{k}: {v}" for k, v in stats.items()] + line = ", ".join(line) + else: + line = "None" + lines += f" {f'({c})'.ljust(maxlen)}{line}\n" + return lines + diff --git a/ngclearn/components/neurons/spiking/LIFCell.py b/ngclearn/components/neurons/spiking/LIFCell.py index 44c2474d..ead51f02 100644 --- a/ngclearn/components/neurons/spiking/LIFCell.py +++ b/ngclearn/components/neurons/spiking/LIFCell.py @@ -116,12 +116,13 @@ class LIFCell(JaxComponent): ## leaky integrate-and-fire cell resist_m: membrane resistance value (Default: 1) thr: base value for adaptive thresholds that govern short-term - plasticity (in milliVolts, or mV) + plasticity (in milliVolts, or mV; default: -52. mV) - v_rest: membrane resting potential (in mV) + v_rest: membrane resting potential (in mV; default: -65 mV) v_reset: membrane reset potential (in mV) -- upon occurrence of a spike, - a neuronal cell's membrane potential will be set to this value + a neuronal cell's membrane potential will be set to this value; + (default: -60 mV) v_decay: decay factor applied to voltage leak (Default: 1.); setting this to 0 mV recovers pure integrate-and-fire (IF) dynamics @@ -131,7 +132,7 @@ class LIFCell(JaxComponent): ## leaky integrate-and-fire cell theta_plus: physical increment to be applied to any threshold value if a spike was emitted - refract_time: relative refractory period time (ms; Default: 1 ms) + refract_time: relative refractory period time (ms; Default: 5 ms) one_spike: if True, a single-spike constraint will be enforced for every time step of neuronal dynamics simulated, i.e., at most, only @@ -146,13 +147,26 @@ class LIFCell(JaxComponent): ## leaky integrate-and-fire cell :Note: setting the integration type to the midpoint method will increase the accuray of the estimate of the cell's evolution at an increase in computational cost (and simulation time) + + surrgoate_type: type of surrogate function to use for approximating a + partial derivative of this cell's spikes w.r.t. its voltage/current + (default: "straight_through") + + :Note: surrogate options available include: "straight_through" + (straight-through estimator), "triangular" (triangular estimator), + "arctan" (arc-tangent estimator), and "secant_lif" (the + LIF-specialized secant estimator) + + lower_clamp_voltage: if True, this will ensure voltage never is below + the value of `v_rest` (default: True) """ @deprecate_args(thr_jitter=None) def __init__(self, name, n_units, tau_m, resist_m=1., thr=-52., v_rest=-65., v_reset=-60., v_decay=1., tau_theta=1e7, theta_plus=0.05, refract_time=5., one_spike=False, integration_type="euler", - surrgoate_type="straight_through", **kwargs): + surrgoate_type="straight_through", lower_clamp_voltage=True, + **kwargs): super().__init__(name, **kwargs) ## Integration properties @@ -163,6 +177,7 @@ def __init__(self, name, n_units, tau_m, resist_m=1., thr=-52., v_rest=-65., self.tau_m = tau_m ## membrane time constant self.resist_m = resist_m ## resistance value self.one_spike = one_spike ## True => constrains system to simulate 1 spike per time step + self.lower_clamp_voltage = lower_clamp_voltage ## True ==> ensures voltage is never < v_rest self.v_rest = v_rest #-65. # mV self.v_reset = v_reset # -60. # -65. # mV (milli-volts) @@ -207,8 +222,8 @@ def __init__(self, name, n_units, tau_m, resist_m=1., thr=-52., v_rest=-65., @staticmethod def _advance_state(t, dt, tau_m, resist_m, v_rest, v_reset, v_decay, refract_T, - thr, tau_theta, theta_plus, one_spike, intgFlag, d_spike_fx, - key, j, v, s, rfr, thr_theta, tols): + thr, tau_theta, theta_plus, one_spike, lower_clamp_voltage, + intgFlag, d_spike_fx, key, j, v, rfr, thr_theta, tols): skey = None ## this is an empty dkey if single_spike mode turned off if one_spike: key, skey = random.split(key, 2) @@ -223,7 +238,9 @@ def _advance_state(t, dt, tau_m, resist_m, v_rest, v_reset, v_decay, refract_T, thr_theta = _update_theta(dt, thr_theta, raw_spikes, tau_theta, theta_plus) ## update tols tols = _update_times(t, s, tols) - return jnp.maximum(v, v_rest), s, raw_spikes, rfr, thr_theta, tols, key, surrogate + if lower_clamp_voltage: ## ensure voltage never < v_rest + v = jnp.maximum(v, v_rest) + return v, s, raw_spikes, rfr, thr_theta, tols, key, surrogate @resolver(_advance_state) def advance_state(self, v, s, s_raw, rfr, thr_theta, tols, key, surrogate): @@ -317,6 +334,7 @@ def help(cls): ## component help function {"v": "Membrane potential/voltage at time t", "rfr": "Current state of (relative) refractory variable", "thr": "Current state of voltage threshold at time t", + "thr_theta": "Current state of homeostatic adaptive threshold at time t", "key": "JAX PRNG key"}, "outputs": {"s": "Emitted spikes/pulses at time t", @@ -331,10 +349,17 @@ def help(cls): ## component help function "v_reset": "Reset membrane potential value", "v_decay": "Voltage leak/decay factor", "tau_theta": "Threshold/homoestatic increment time constant", - "theta_plus": "Amount to increment threshold by upon occurrence of spike", + "theta_plus": "Amount to increment threshold by upon occurrence " + "of spike", "refract_time": "Length of relative refractory period (ms)", - "one_spike": "Should only one spike be sampled/allowed to emit at any given time step?", - "integration_type": "Type of numerical integration to use for the cell dynamics" + "one_spike": "Should only one spike be sampled/allowed to emit at " + "any given time step?", + "integration_type": "Type of numerical integration to use for the " + "cell dynamics", + "surrgoate_type": "Type of surrogate function to use approximate " + "derivative of spike w.r.t. voltage/current", + "lower_bound_clamp": "Should voltage be lower bounded to be never " + "be below `v_rest`" } info = {cls.__name__: properties, "compartments": compartment_props, diff --git a/ngclearn/components/neurons/spiking/__init__.py b/ngclearn/components/neurons/spiking/__init__.py index cd9aa181..2934eda9 100644 --- a/ngclearn/components/neurons/spiking/__init__.py +++ b/ngclearn/components/neurons/spiking/__init__.py @@ -1,6 +1,7 @@ ## point to standard spiking cell component types from .sLIFCell import SLIFCell from .LIFCell import LIFCell +from .IFCell import IFCell from .WTASCell import WTASCell from .quadLIFCell import QuadLIFCell from .adExCell import AdExCell From 27a61ef4a41ef2b43e3aca1b0e8deb6228eb37da Mon Sep 17 00:00:00 2001 From: ago109 Date: Wed, 24 Jul 2024 16:59:30 -0400 Subject: [PATCH 04/61] mod to latency-cell --- ngclearn/components/input_encoders/latencyCell.py | 1 + 1 file changed, 1 insertion(+) diff --git a/ngclearn/components/input_encoders/latencyCell.py b/ngclearn/components/input_encoders/latencyCell.py index 1acf199f..b3666bcb 100755 --- a/ngclearn/components/input_encoders/latencyCell.py +++ b/ngclearn/components/input_encoders/latencyCell.py @@ -4,6 +4,7 @@ from ngclearn.utils.model_utils import clamp_min, clamp_max from jax import numpy as jnp, random, jit from functools import partial +from ngcsimlib.logger import info @jit def _update_times(t, s, tols): From 05a97f029f9246015f4c6058fb1b0dadb8095050 Mon Sep 17 00:00:00 2001 From: Will Gebhardt Date: Thu, 25 Jul 2024 10:04:47 -0400 Subject: [PATCH 05/61] updated the poissonCell to be a true poisson --- .../components/input_encoders/poissonCell.py | 125 +++++++++--------- 1 file changed, 61 insertions(+), 64 deletions(-) diff --git a/ngclearn/components/input_encoders/poissonCell.py b/ngclearn/components/input_encoders/poissonCell.py index 115afbc9..f27918b8 100644 --- a/ngclearn/components/input_encoders/poissonCell.py +++ b/ngclearn/components/input_encoders/poissonCell.py @@ -1,52 +1,15 @@ from ngclearn import resolver, Component, Compartment from ngclearn.components.jaxComponent import JaxComponent from ngclearn.utils import tensorstats -from jax import numpy as jnp, random, jit +from jax import numpy as jnp, random, jit, scipy from functools import partial +from ngcsimlib.deprecators import deprecate_args -@jit -def _update_times(t, s, tols): - """ - Updates time-of-last-spike (tols) variable. - - Args: - t: current time (a scalar/int value) - - s: binary spike vector - - tols: current time-of-last-spike variable - - Returns: - updated tols variable - """ - _tols = (1. - s) * tols + (s * t) - return _tols - -@partial(jit, static_argnums=[3]) -def _sample_poisson(dkey, data, dt, fmax=63.75): - """ - Samples a Poisson spike train on-the-fly. - - Args: - dkey: JAX key to drive stochasticity/noise - - data: sensory data (vector/matrix) - - dt: integration time constant - - fmax: maximum frequency (Hz) - - Returns: - binary spikes - """ - pspike = data * (dt/1000.) * fmax - eps = random.uniform(dkey, data.shape, minval=0., maxval=1., dtype=jnp.float32) - s_t = (eps < pspike).astype(jnp.float32) - return s_t class PoissonCell(JaxComponent): """ - A Poisson cell that produces approximately Poisson-distributed spikes on-the-fly. + A Poisson cell that produces approximately Poisson-distributed spikes + on-the-fly. | --- Cell Input Compartments: --- | inputs - input (takes in external signals) @@ -61,49 +24,78 @@ class PoissonCell(JaxComponent): n_units: number of cellular entities (neural population size) - max_freq: maximum frequency (in Hertz) of this Poisson spike train (must be > 0.) + max_freq: maximum frequency (in Hertz) of this Poisson spike train ( + must be > 0.) """ # Define Functions - def __init__(self, name, n_units, max_freq=63.75, batch_size=1, **kwargs): + @deprecate_args(target_freq="max_freq") + def __init__(self, name, n_units, target_freq=63.75, batch_size=1, + **kwargs): super().__init__(name, **kwargs) ## Poisson meta-parameters - self.max_freq = max_freq ## maximum frequency (in Hertz/Hz) + self.target_freq = target_freq ## maximum frequency (in Hertz/Hz) ## Layer Size Setup self.batch_size = batch_size self.n_units = n_units + _key, subkey = random.split(self.key.value, 2) + self.key.set(_key) ## Compartment setup restVals = jnp.zeros((self.batch_size, self.n_units)) - self.inputs = Compartment(restVals, display_name="Input Stimulus") # input compartment - self.outputs = Compartment(restVals, display_name="Spikes") # output compartment - self.tols = Compartment(restVals, display_name="Time-of-Last-Spike", units="ms") # time of last spike + self.inputs = Compartment(restVals, + display_name="Input Stimulus") # input + # compartment + self.outputs = Compartment(restVals, + display_name="Spikes") # output compartment + self.tols = Compartment(restVals, display_name="Time-of-Last-Spike", + units="ms") # time of last spike + self.targets = Compartment( + random.uniform(subkey, (self.batch_size, self.n_units), minval=0., + maxval=1.)) @staticmethod - def _advance_state(t, dt, max_freq, key, inputs, tols): - key, *subkeys = random.split(key, 2) - outputs = _sample_poisson(subkeys[0], data=inputs, dt=dt, fmax=max_freq) - tols = _update_times(t, outputs, tols) - return outputs, tols, key + def _advance_state(t, dt, target_freq, key, inputs, targets, tols): + ms_per_second = 1000 # ms/s + events_per_ms = target_freq / ms_per_second # e/s s/ms -> e/ms + ms_per_event = 1 / events_per_ms # ms/e + time_step_per_event = ms_per_event / dt # ms/e * ts/ms -> ts / e + + cdf = scipy.special.gammaincc((t + dt) - tols, + time_step_per_event/inputs) + outputs = (targets < cdf).astype(jnp.float32) + + key, subkey = random.split(key, 2) + targets = (targets * (1 - outputs) + random.uniform(subkey, + targets.shape) * + outputs) + + tols = tols * (1. - outputs) + t * outputs + return outputs, tols, key, targets @resolver(_advance_state) - def advance_state(self, outputs, tols, key): + def advance_state(self, outputs, tols, key, targets): self.outputs.set(outputs) self.tols.set(tols) self.key.set(key) + self.targets.set(targets) @staticmethod - def _reset(batch_size, n_units): + def _reset(batch_size, n_units, key): restVals = jnp.zeros((batch_size, n_units)) - return restVals, restVals, restVals + key, subkey = random.split(key, 2) + targets = random.uniform(subkey, (batch_size, n_units)) + return restVals, restVals, restVals, targets, key @resolver(_reset) - def reset(self, inputs, outputs, tols): + def reset(self, inputs, outputs, tols, targets, key): self.inputs.set(inputs) self.outputs.set(outputs) self.tols.set(tols) + self.key.set(key) + self.targets.set(targets) def save(self, directory, **kwargs): file_name = directory + "/" + self.name + ".npz" @@ -115,19 +107,21 @@ def load(self, directory, **kwargs): self.key.set(data['key']) @classmethod - def help(cls): ## component help function + def help(cls): ## component help function properties = { "cell_type": "PoissonCell - samples input to produce spikes, " - "where dimension is a probability proportional to " - "the dimension's magnitude/value/intensity and " - "constrained by a maximum spike frequency (spikes follow " + "where dimension is a probability proportional to " + "the dimension's magnitude/value/intensity and " + "constrained by a maximum spike frequency (spikes " + "follow " "a Poisson distribution)" } compartment_props = { "inputs": {"inputs": "Takes in external input signal values"}, "states": - {"key": "JAX PRNG key"}, + {"key": "JAX PRNG key", + "targets": "Target cdf for the Poisson distribution"}, "outputs": {"tols": "Time-of-last-spike", "outputs": "Binary spike values emitted at time t"}, @@ -135,16 +129,17 @@ def help(cls): ## component help function hyperparams = { "n_units": "Number of neuronal cells to model in this layer", "batch_size": "Batch size dimension of this component", - "max_freq": "Maximum spike frequency of the train produced", + "target_freq": "Maximum spike frequency of the train produced", } info = {cls.__name__: properties, "compartments": compartment_props, - "dynamics": "~ Poisson(x; max_freq)", + "dynamics": "~ Poisson(x; target_freq)", "hyperparameters": hyperparams} return info def __repr__(self): - comps = [varname for varname in dir(self) if Compartment.is_compartment(getattr(self, varname))] + comps = [varname for varname in dir(self) if + Compartment.is_compartment(getattr(self, varname))] maxlen = max(len(c) for c in comps) + 5 lines = f"[{self.__class__.__name__}] PATH: {self.name}\n" for c in comps: @@ -157,8 +152,10 @@ def __repr__(self): lines += f" {f'({c})'.ljust(maxlen)}{line}\n" return lines + if __name__ == '__main__': from ngcsimlib.context import Context + with Context("Bar") as bar: X = PoissonCell("X", 9) print(X) From efa61a532a87ae2831af7935700f28870d3da762 Mon Sep 17 00:00:00 2001 From: ago109 Date: Thu, 25 Jul 2024 11:40:26 -0400 Subject: [PATCH 06/61] fixed minor bug in deprecation for poiss/bern --- .../components/input_encoders/bernoulliCell.py | 2 +- .../components/input_encoders/poissonCell.py | 18 ++++++++++++++++-- 2 files changed, 17 insertions(+), 3 deletions(-) diff --git a/ngclearn/components/input_encoders/bernoulliCell.py b/ngclearn/components/input_encoders/bernoulliCell.py index fa036302..841804e9 100755 --- a/ngclearn/components/input_encoders/bernoulliCell.py +++ b/ngclearn/components/input_encoders/bernoulliCell.py @@ -84,7 +84,7 @@ class BernoulliCell(JaxComponent): target_freq: maximum frequency (in Hertz) of this Bernoulli spike train (must be > 0.) """ - @deprecate_args(target_freq="max_freq") + @deprecate_args(max_freq="target_freq") def __init__(self, name, n_units, target_freq=63.75, batch_size=1, **kwargs): super().__init__(name, **kwargs) diff --git a/ngclearn/components/input_encoders/poissonCell.py b/ngclearn/components/input_encoders/poissonCell.py index f27918b8..36b968a8 100644 --- a/ngclearn/components/input_encoders/poissonCell.py +++ b/ngclearn/components/input_encoders/poissonCell.py @@ -4,7 +4,7 @@ from jax import numpy as jnp, random, jit, scipy from functools import partial from ngcsimlib.deprecators import deprecate_args - +from ngcsimlib.logger import info, warn class PoissonCell(JaxComponent): """ @@ -29,7 +29,7 @@ class PoissonCell(JaxComponent): """ # Define Functions - @deprecate_args(target_freq="max_freq") + @deprecate_args(max_freq="target_freq") def __init__(self, name, n_units, target_freq=63.75, batch_size=1, **kwargs): super().__init__(name, **kwargs) @@ -56,6 +56,20 @@ def __init__(self, name, n_units, target_freq=63.75, batch_size=1, random.uniform(subkey, (self.batch_size, self.n_units), minval=0., maxval=1.)) + def validate(self, dt, **validation_kwargs): + ## check for unstable combinations of dt and target-frequency meta-params + valid = super().validate(**validation_kwargs) + events_per_timestep = (dt/1000.) * self.target_freq ## compute scaled probability + if events_per_timestep > 1.: + valid = False + warn( + f"{self.name} will be unable to make as many temporal events as " + f"requested! ({events_per_timestep} events/timestep) Unstable " + f"combination of dt = {dt} and target_freq = {self.target_freq} " + f"being used!" + ) + return valid + @staticmethod def _advance_state(t, dt, target_freq, key, inputs, targets, tols): ms_per_second = 1000 # ms/s From 223d3c0244c746046716acb241960053e9c4bb70 Mon Sep 17 00:00:00 2001 From: ago109 Date: Thu, 25 Jul 2024 11:43:24 -0400 Subject: [PATCH 07/61] fixed minor bug in deprecation for poiss/bern --- ngclearn/components/input_encoders/bernoulliCell.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ngclearn/components/input_encoders/bernoulliCell.py b/ngclearn/components/input_encoders/bernoulliCell.py index 841804e9..ea3b26cc 100755 --- a/ngclearn/components/input_encoders/bernoulliCell.py +++ b/ngclearn/components/input_encoders/bernoulliCell.py @@ -85,7 +85,7 @@ class BernoulliCell(JaxComponent): """ @deprecate_args(max_freq="target_freq") - def __init__(self, name, n_units, target_freq=63.75, batch_size=1, **kwargs): + def __init__(self, name, n_units, target_freq=0., batch_size=1, **kwargs): super().__init__(name, **kwargs) ## Constrained Bernoulli meta-parameters From 9afaadfe65df2bd2597fb49d75666efc58eb3c31 Mon Sep 17 00:00:00 2001 From: ago109 Date: Thu, 25 Jul 2024 11:55:08 -0400 Subject: [PATCH 08/61] fixed validation fun in bern/poiss --- ngclearn/components/input_encoders/bernoulliCell.py | 7 +++++-- ngclearn/components/input_encoders/poissonCell.py | 9 ++++++--- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/ngclearn/components/input_encoders/bernoulliCell.py b/ngclearn/components/input_encoders/bernoulliCell.py index ea3b26cc..b061b431 100755 --- a/ngclearn/components/input_encoders/bernoulliCell.py +++ b/ngclearn/components/input_encoders/bernoulliCell.py @@ -101,9 +101,12 @@ def __init__(self, name, n_units, target_freq=0., batch_size=1, **kwargs): self.outputs = Compartment(restVals, display_name="Spikes") # output compartment self.tols = Compartment(restVals, display_name="Time-of-Last-Spike", units="ms") # time of last spike - def validate(self, dt, **validation_kwargs): - ## check for unstable combinations of dt and target-frequency meta-params + def validate(self, dt=None, **validation_kwargs): valid = super().validate(**validation_kwargs) + if dt is None: + warn(f"{self.name} requires a validation kwarg of `dt`") + return False + ## check for unstable combinations of dt and target-frequency meta-params events_per_timestep = (dt/1000.) * self.target_freq ## compute scaled probability if events_per_timestep > 1.: valid = False diff --git a/ngclearn/components/input_encoders/poissonCell.py b/ngclearn/components/input_encoders/poissonCell.py index 36b968a8..8c5fba99 100644 --- a/ngclearn/components/input_encoders/poissonCell.py +++ b/ngclearn/components/input_encoders/poissonCell.py @@ -56,10 +56,13 @@ def __init__(self, name, n_units, target_freq=63.75, batch_size=1, random.uniform(subkey, (self.batch_size, self.n_units), minval=0., maxval=1.)) - def validate(self, dt, **validation_kwargs): - ## check for unstable combinations of dt and target-frequency meta-params + def validate(self, dt=None, **validation_kwargs): valid = super().validate(**validation_kwargs) - events_per_timestep = (dt/1000.) * self.target_freq ## compute scaled probability + if dt is None: + warn(f"{self.name} requires a validation kwarg of `dt`") + return False + ## check for unstable combinations of dt and target-frequency meta-params + events_per_timestep = (dt / 1000.) * self.target_freq ## compute scaled probability if events_per_timestep > 1.: valid = False warn( From bf720945e170a85008a72dde92d0a2f9835516ca Mon Sep 17 00:00:00 2001 From: ago109 Date: Thu, 25 Jul 2024 15:51:08 -0400 Subject: [PATCH 09/61] moved back and cleaned up bernoulli and poisson cells --- .../input_encoders/bernoulliCell.py | 77 +-------------- .../components/input_encoders/poissonCell.py | 97 +++++++++---------- 2 files changed, 51 insertions(+), 123 deletions(-) diff --git a/ngclearn/components/input_encoders/bernoulliCell.py b/ngclearn/components/input_encoders/bernoulliCell.py index b061b431..b0afda29 100755 --- a/ngclearn/components/input_encoders/bernoulliCell.py +++ b/ngclearn/components/input_encoders/bernoulliCell.py @@ -24,49 +24,10 @@ def _update_times(t, s, tols): _tols = (1. - s) * tols + (s * t) return _tols -@jit -def _sample_bernoulli(dkey, data): - """ - Samples a Bernoulli spike train on-the-fly - - Args: - dkey: JAX key to drive stochasticity/noise - - data: sensory data (vector/matrix) - - Returns: - binary spikes - """ - s_t = random.bernoulli(dkey, p=data).astype(jnp.float32) - return s_t - -@partial(jit, static_argnums=[3]) -def _sample_constrained_bernoulli(dkey, data, dt, fmax=63.75): - """ - Samples a Bernoulli spike train on-the-fly that is constrained to emit - at a particular rate over a time window. - - Args: - dkey: JAX key to drive stochasticity/noise - - data: sensory data (vector/matrix) - - dt: integration time constant - - fmax: maximum frequency (Hz) - - Returns: - binary spikes - """ - pspike = data * (dt/1000.) * fmax - eps = random.uniform(dkey, data.shape, minval=0., maxval=1., dtype=jnp.float32) - s_t = (eps < pspike).astype(jnp.float32) - return s_t - class BernoulliCell(JaxComponent): """ - A Bernoulli cell that produces variations of Bernoulli-distributed spikes - on-the-fly (including constrained-rate trains). + A Bernoulli cell that produces spikes by sampling a Bernoulli distribution + on-the-fly (to produce data-scaled Bernoulli spike trains). | --- Cell Input Compartments: --- | inputs - input (takes in external signals) @@ -80,17 +41,11 @@ class BernoulliCell(JaxComponent): name: the string name of this cell n_units: number of cellular entities (neural population size) - - target_freq: maximum frequency (in Hertz) of this Bernoulli spike train (must be > 0.) """ - @deprecate_args(max_freq="target_freq") - def __init__(self, name, n_units, target_freq=0., batch_size=1, **kwargs): + def __init__(self, name, n_units, batch_size=1, **kwargs): super().__init__(name, **kwargs) - ## Constrained Bernoulli meta-parameters - self.target_freq = target_freq ## maximum frequency (in Hertz/Hz) - ## Layer Size Setup self.batch_size = batch_size self.n_units = n_units @@ -101,32 +56,10 @@ def __init__(self, name, n_units, target_freq=0., batch_size=1, **kwargs): self.outputs = Compartment(restVals, display_name="Spikes") # output compartment self.tols = Compartment(restVals, display_name="Time-of-Last-Spike", units="ms") # time of last spike - def validate(self, dt=None, **validation_kwargs): - valid = super().validate(**validation_kwargs) - if dt is None: - warn(f"{self.name} requires a validation kwarg of `dt`") - return False - ## check for unstable combinations of dt and target-frequency meta-params - events_per_timestep = (dt/1000.) * self.target_freq ## compute scaled probability - if events_per_timestep > 1.: - valid = False - warn( - f"{self.name} will be unable to make as many temporal events as " - f"requested! ({events_per_timestep} events/timestep) Unstable " - f"combination of dt = {dt} and target_freq = {self.target_freq} " - f"being used!" - ) - return valid - @staticmethod - def _advance_state(t, dt, target_freq, key, inputs, tols): + def _advance_state(t, key, inputs, tols): key, *subkeys = random.split(key, 2) - if target_freq > 0.: - outputs = _sample_constrained_bernoulli( ## sample Bernoulli w/ target rate - subkeys[0], data=inputs, dt=dt, fmax=target_freq - ) - else: - outputs = _sample_bernoulli(subkeys[0], data=inputs) + outputs = random.bernoulli(subkeys[0], p=inputs).astype(jnp.float32) tols = _update_times(t, outputs, tols) return outputs, tols, key diff --git a/ngclearn/components/input_encoders/poissonCell.py b/ngclearn/components/input_encoders/poissonCell.py index 8c5fba99..3c97fdb1 100644 --- a/ngclearn/components/input_encoders/poissonCell.py +++ b/ngclearn/components/input_encoders/poissonCell.py @@ -1,15 +1,33 @@ from ngclearn import resolver, Component, Compartment from ngclearn.components.jaxComponent import JaxComponent +from jax import numpy as jnp, random, jit from ngclearn.utils import tensorstats -from jax import numpy as jnp, random, jit, scipy from functools import partial from ngcsimlib.deprecators import deprecate_args from ngcsimlib.logger import info, warn +@jit +def _update_times(t, s, tols): + """ + Updates time-of-last-spike (tols) variable. + + Args: + t: current time (a scalar/int value) + + s: binary spike vector + + tols: current time-of-last-spike variable + + Returns: + updated tols variable + """ + _tols = (1. - s) * tols + (s * t) + return _tols + class PoissonCell(JaxComponent): """ - A Poisson cell that produces approximately Poisson-distributed spikes - on-the-fly. + A Poisson cell that samples a homogeneous Poisson process on-the-fly to + produce a spike train. | --- Cell Input Compartments: --- | inputs - input (takes in external signals) @@ -24,37 +42,25 @@ class PoissonCell(JaxComponent): n_units: number of cellular entities (neural population size) - max_freq: maximum frequency (in Hertz) of this Poisson spike train ( - must be > 0.) + target_freq: maximum frequency (in Hertz) of this Bernoulli spike train (must be > 0.) """ - # Define Functions @deprecate_args(max_freq="target_freq") - def __init__(self, name, n_units, target_freq=63.75, batch_size=1, - **kwargs): + def __init__(self, name, n_units, target_freq=0., batch_size=1, **kwargs): super().__init__(name, **kwargs) - ## Poisson meta-parameters + ## Constrained Bernoulli meta-parameters self.target_freq = target_freq ## maximum frequency (in Hertz/Hz) ## Layer Size Setup self.batch_size = batch_size self.n_units = n_units - _key, subkey = random.split(self.key.value, 2) - self.key.set(_key) - ## Compartment setup + # Compartments (state of the cell, parameters, will be updated through stateless calls) restVals = jnp.zeros((self.batch_size, self.n_units)) - self.inputs = Compartment(restVals, - display_name="Input Stimulus") # input - # compartment - self.outputs = Compartment(restVals, - display_name="Spikes") # output compartment - self.tols = Compartment(restVals, display_name="Time-of-Last-Spike", - units="ms") # time of last spike - self.targets = Compartment( - random.uniform(subkey, (self.batch_size, self.n_units), minval=0., - maxval=1.)) + self.inputs = Compartment(restVals, display_name="Input Stimulus") # input compartment + self.outputs = Compartment(restVals, display_name="Spikes") # output compartment + self.tols = Compartment(restVals, display_name="Time-of-Last-Spike", units="ms") # time of last spike def validate(self, dt=None, **validation_kwargs): valid = super().validate(**validation_kwargs) @@ -62,7 +68,7 @@ def validate(self, dt=None, **validation_kwargs): warn(f"{self.name} requires a validation kwarg of `dt`") return False ## check for unstable combinations of dt and target-frequency meta-params - events_per_timestep = (dt / 1000.) * self.target_freq ## compute scaled probability + events_per_timestep = (dt/1000.) * self.target_freq ## compute scaled probability if events_per_timestep > 1.: valid = False warn( @@ -74,54 +80,43 @@ def validate(self, dt=None, **validation_kwargs): return valid @staticmethod - def _advance_state(t, dt, target_freq, key, inputs, targets, tols): - ms_per_second = 1000 # ms/s - events_per_ms = target_freq / ms_per_second # e/s s/ms -> e/ms - ms_per_event = 1 / events_per_ms # ms/e - time_step_per_event = ms_per_event / dt # ms/e * ts/ms -> ts / e - - cdf = scipy.special.gammaincc((t + dt) - tols, - time_step_per_event/inputs) - outputs = (targets < cdf).astype(jnp.float32) - - key, subkey = random.split(key, 2) - targets = (targets * (1 - outputs) + random.uniform(subkey, - targets.shape) * - outputs) - - tols = tols * (1. - outputs) + t * outputs - return outputs, tols, key, targets + def _advance_state(t, dt, target_freq, key, inputs, tols): + key, *subkeys = random.split(key, 2) + pspike = inputs * (dt / 1000.) * target_freq + eps = random.uniform(subkeys[0], inputs.shape, minval=0., maxval=1., + dtype=jnp.float32) + outputs = (eps < pspike).astype(jnp.float32) + tols = _update_times(t, outputs, tols) + return outputs, tols, key @resolver(_advance_state) - def advance_state(self, outputs, tols, key, targets): + def advance_state(self, outputs, tols, key): self.outputs.set(outputs) self.tols.set(tols) self.key.set(key) - self.targets.set(targets) @staticmethod - def _reset(batch_size, n_units, key): + def _reset(batch_size, n_units): restVals = jnp.zeros((batch_size, n_units)) - key, subkey = random.split(key, 2) - targets = random.uniform(subkey, (batch_size, n_units)) - return restVals, restVals, restVals, targets, key + return restVals, restVals, restVals @resolver(_reset) - def reset(self, inputs, outputs, tols, targets, key): + def reset(self, inputs, outputs, tols): self.inputs.set(inputs) - self.outputs.set(outputs) + self.outputs.set(outputs) #None self.tols.set(tols) - self.key.set(key) - self.targets.set(targets) def save(self, directory, **kwargs): + target_freq = (self.target_freq if isinstance(self.target_freq, float) + else jnp.ones([[self.target_freq]])) file_name = directory + "/" + self.name + ".npz" - jnp.savez(file_name, key=self.key.value) + jnp.savez(file_name, key=self.key.value, target_freq=target_freq) def load(self, directory, **kwargs): file_name = directory + "/" + self.name + ".npz" data = jnp.load(file_name) self.key.set(data['key']) + self.target_freq = data['target_freq'] @classmethod def help(cls): ## component help function From c894b8af3e2e073819007052d26525784a1344d7 Mon Sep 17 00:00:00 2001 From: ago109 Date: Thu, 25 Jul 2024 17:21:11 -0400 Subject: [PATCH 10/61] added threshold-clipping to latency cell --- .../components/input_encoders/latencyCell.py | 30 +++++++++++++------ 1 file changed, 21 insertions(+), 9 deletions(-) diff --git a/ngclearn/components/input_encoders/latencyCell.py b/ngclearn/components/input_encoders/latencyCell.py index b3666bcb..cf6db5ca 100755 --- a/ngclearn/components/input_encoders/latencyCell.py +++ b/ngclearn/components/input_encoders/latencyCell.py @@ -149,13 +149,16 @@ class LatencyCell(JaxComponent): :Note: if this set to True, you will need to choose a useful value for the "num_steps" argument (>1), depending on how many steps simulated + clip_spikes: should values under threshold be removed/suppressed? + (default: False) + num_steps: number of discrete time steps to consider for normalized latency code (only useful if "normalize" is set to True) (Default: 1) """ # Define Functions def __init__(self, name, n_units, tau=1., threshold=0.01, first_spike_time=0., - linearize=False, normalize=False, num_steps=1., + linearize=False, normalize=False, clip_spikes=False, num_steps=1., batch_size=1, **kwargs): super().__init__(name, **kwargs) @@ -164,6 +167,7 @@ def __init__(self, name, n_units, tau=1., threshold=0.01, first_spike_time=0., self.tau = tau self.threshold = threshold self.linearize = linearize + self.clip_spikes = clip_spikes ## normalize latency code s.t. final spike(s) occur w/in num_steps self.normalize = normalize self.num_steps = num_steps @@ -176,17 +180,22 @@ def __init__(self, name, n_units, tau=1., threshold=0.01, first_spike_time=0., restVals = jnp.zeros((self.batch_size, self.n_units)) self.inputs = Compartment(restVals, display_name="Input Stimulus") # input compartment self.outputs = Compartment(restVals, display_name="Spikes") # output compartment - self.mask = Compartment(restVals, display_name="Mask Variable") # output compartment + self.mask = Compartment(restVals, display_name="Spike Time Mask") + self.clip_mask = Compartment(restVals, display_name="Clip Mask") self.tols = Compartment(restVals, display_name="Time-of-Last-Spike", units="ms") # time of last spike self.targ_sp_times = Compartment(restVals, display_name="Target Spike Time", units="ms") #self.reset() @staticmethod def _calc_spike_times(linearize, tau, threshold, first_spike_time, num_steps, - normalize, inputs): + normalize, clip_spikes, inputs): ## would call this function before processing a spike train (at start) data = inputs - if linearize == True: ## linearize spike time calculation + if clip_spikes: + clip_mask = (data < threshold) * 1. ## find values under threshold + else: + clip_mask = data * 0. ## all values allowed to fire spikes + if linearize: ## linearize spike time calculation stimes = _calc_spike_times_linear(data, tau, threshold, first_spike_time, num_steps, normalize) @@ -197,18 +206,20 @@ def _calc_spike_times(linearize, tau, threshold, first_spike_time, num_steps, num_steps=num_steps, normalize=normalize) targ_sp_times = stimes #* calcEvent + targ_sp_times * (1. - calcEvent) - return targ_sp_times + return targ_sp_times, clip_mask @resolver(_calc_spike_times) - def calc_spike_times(self, targ_sp_times): + def calc_spike_times(self, targ_sp_times, clip_mask): self.targ_sp_times.set(targ_sp_times) + self.clip_mask.set(clip_mask) @staticmethod - def _advance_state(t, dt, key, inputs, mask, targ_sp_times, tols): + def _advance_state(t, dt, key, inputs, mask, clip_mask, targ_sp_times, tols): key, *subkeys = random.split(key, 2) data = inputs ## get sensory pattern data / features spikes, spk_mask = _extract_spike(targ_sp_times, t, mask) ## get spikes at t tols = _update_times(t, spikes, tols) + spikes = spikes * (1. - clip_mask) return spikes, tols, spk_mask, targ_sp_times, key @resolver(_advance_state) @@ -222,14 +233,15 @@ def advance_state(self, outputs, tols, mask, targ_sp_times, key): @staticmethod def _reset(batch_size, n_units): restVals = jnp.zeros((batch_size, n_units)) - return (restVals, restVals, restVals, restVals, restVals) + return (restVals, restVals, restVals, restVals, restVals, restVals) @resolver(_reset) - def reset(self, inputs, outputs, tols, mask, targ_sp_times): + def reset(self, inputs, outputs, tols, mask, clip_mask, targ_sp_times): self.inputs.set(inputs) self.outputs.set(outputs) self.tols.set(tols) self.mask.set(mask) + self.clip_mask.set(clip_mask) self.targ_sp_times.set(targ_sp_times) def save(self, directory, **kwargs): From ba08453b09f6a90954218fa0a02272d92ce5192e Mon Sep 17 00:00:00 2001 From: ago109 Date: Fri, 26 Jul 2024 13:33:07 -0400 Subject: [PATCH 11/61] updates to if/lif --- ngclearn/components/neurons/spiking/IFCell.py | 16 ++++++++-------- ngclearn/components/neurons/spiking/LIFCell.py | 16 ++++++++-------- 2 files changed, 16 insertions(+), 16 deletions(-) diff --git a/ngclearn/components/neurons/spiking/IFCell.py b/ngclearn/components/neurons/spiking/IFCell.py index 2c9acf52..68e51bdf 100755 --- a/ngclearn/components/neurons/spiking/IFCell.py +++ b/ngclearn/components/neurons/spiking/IFCell.py @@ -213,21 +213,21 @@ def reset(self, j, v, s, rfr, tols, surrogate): def save(self, directory, **kwargs): ## do a protected save of constants, depending on whether they are floats or arrays tau_m = (self.tau_m if isinstance(self.tau_m, float) - else jnp.ones([[self.tau_m]])) + else jnp.asarray([[self.tau_m * 1.]])) thr = (self.thr if isinstance(self.thr, float) - else jnp.ones([[self.thr]])) + else jnp.asarray([[self.thr * 1.]])) v_rest = (self.v_rest if isinstance(self.v_rest, float) - else jnp.ones([[self.v_rest]])) + else jnp.asarray([[self.v_rest * 1.]])) v_reset = (self.v_reset if isinstance(self.v_reset, float) - else jnp.ones([[self.v_reset]])) + else jnp.asarray([[self.v_reset * 1.]])) v_decay = (self.v_decay if isinstance(self.v_decay, float) - else jnp.ones([[self.v_decay]])) + else jnp.asarray([[self.v_decay * 1.]])) resist_m = (self.resist_m if isinstance(self.resist_m, float) - else jnp.ones([[self.resist_m]])) + else jnp.asarray([[self.resist_m * 1.]])) tau_theta = (self.tau_theta if isinstance(self.tau_theta, float) - else jnp.ones([[self.tau_theta]])) + else jnp.asarray([[self.tau_theta * 1.]])) theta_plus = (self.theta_plus if isinstance(self.theta_plus, float) - else jnp.ones([[self.theta_plus]])) + else jnp.asarray([[self.theta_plus * 1.]])) file_name = directory + "/" + self.name + ".npz" jnp.savez(file_name, diff --git a/ngclearn/components/neurons/spiking/LIFCell.py b/ngclearn/components/neurons/spiking/LIFCell.py index ead51f02..d8a0d763 100644 --- a/ngclearn/components/neurons/spiking/LIFCell.py +++ b/ngclearn/components/neurons/spiking/LIFCell.py @@ -279,21 +279,21 @@ def reset(self, j, v, s, s_raw, rfr, tols, surrogate): def save(self, directory, **kwargs): ## do a protected save of constants, depending on whether they are floats or arrays tau_m = (self.tau_m if isinstance(self.tau_m, float) - else jnp.ones([[self.tau_m]])) + else jnp.asarray([[self.tau_m * 1.]])) thr = (self.thr if isinstance(self.thr, float) - else jnp.ones([[self.thr]])) + else jnp.asarray([[self.thr * 1.]])) v_rest = (self.v_rest if isinstance(self.v_rest, float) - else jnp.ones([[self.v_rest]])) + else jnp.asarray([[self.v_rest * 1.]])) v_reset = (self.v_reset if isinstance(self.v_reset, float) - else jnp.ones([[self.v_reset]])) + else jnp.asarray([[self.v_reset * 1.]])) v_decay = (self.v_decay if isinstance(self.v_decay, float) - else jnp.ones([[self.v_decay]])) + else jnp.asarray([[self.v_decay * 1.]])) resist_m = (self.resist_m if isinstance(self.resist_m, float) - else jnp.ones([[self.resist_m]])) + else jnp.asarray([[self.resist_m * 1.]])) tau_theta = (self.tau_theta if isinstance(self.tau_theta, float) - else jnp.ones([[self.tau_theta]])) + else jnp.asarray([[self.tau_theta * 1.]])) theta_plus = (self.theta_plus if isinstance(self.theta_plus, float) - else jnp.ones([[self.theta_plus]])) + else jnp.asarray([[self.theta_plus * 1.]])) file_name = directory + "/" + self.name + ".npz" jnp.savez(file_name, From 9c932b182d65cd88915aafc2ee3c3a657ef30c9e Mon Sep 17 00:00:00 2001 From: ago109 Date: Fri, 26 Jul 2024 16:46:57 -0400 Subject: [PATCH 12/61] added batch-size arg to slif --- ngclearn/components/neurons/spiking/sLIFCell.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/ngclearn/components/neurons/spiking/sLIFCell.py b/ngclearn/components/neurons/spiking/sLIFCell.py index 37286f3f..e65adae5 100644 --- a/ngclearn/components/neurons/spiking/sLIFCell.py +++ b/ngclearn/components/neurons/spiking/sLIFCell.py @@ -211,7 +211,8 @@ class SLIFCell(JaxComponent): ## leaky integrate-and-fire cell # Define Functions def __init__(self, name, n_units, tau_m, resist_m, thr, resist_inh=0., thr_persist=False, thr_gain=0.0, thr_leak=0.0, rho_b=0., - refract_time=0., sticky_spikes=False, thr_jitter=0.05, **kwargs): + refract_time=0., sticky_spikes=False, thr_jitter=0.05, + batch_size=1, **kwargs): super().__init__(name, **kwargs) ## membrane parameter setup (affects ODE integration) @@ -233,7 +234,7 @@ def __init__(self, name, n_units, tau_m, resist_m, thr, resist_inh=0., ## Layer Size Setup self.n_units = n_units - self.batch_size = 1 + self.batch_size = batch_size ## Adaptive threshold setup self.rho_b = rho_b From 03940e9136f90491448cbc21b96df146834ff2d4 Mon Sep 17 00:00:00 2001 From: ago109 Date: Sat, 27 Jul 2024 14:02:04 -0400 Subject: [PATCH 13/61] fixed minor load bug in lif-cell --- ngclearn/components/neurons/spiking/LIFCell.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ngclearn/components/neurons/spiking/LIFCell.py b/ngclearn/components/neurons/spiking/LIFCell.py index d8a0d763..af6c6683 100644 --- a/ngclearn/components/neurons/spiking/LIFCell.py +++ b/ngclearn/components/neurons/spiking/LIFCell.py @@ -307,7 +307,7 @@ def save(self, directory, **kwargs): def load(self, directory, seeded=False, **kwargs): file_name = directory + "/" + self.name + ".npz" data = jnp.load(file_name) - self.thr_theta.set(data['thr_theta']) + self.thr_theta.set(data['threshold_theta']) ## constants loaded in self.tau_m = data['tau_m'] self.thr = data['thr'] From 6bc5cd8585d6d8f956bcd58e1bfce49b43974738 Mon Sep 17 00:00:00 2001 From: Alexander Ororbia Date: Sat, 27 Jul 2024 14:06:23 -0400 Subject: [PATCH 14/61] fixed a blocking jit-partial call in lif update_theta method; when loading --- ngclearn/components/neurons/spiking/LIFCell.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ngclearn/components/neurons/spiking/LIFCell.py b/ngclearn/components/neurons/spiking/LIFCell.py index af6c6683..ff6edc89 100644 --- a/ngclearn/components/neurons/spiking/LIFCell.py +++ b/ngclearn/components/neurons/spiking/LIFCell.py @@ -72,7 +72,7 @@ def _run_cell(dt, j, v, v_thr, v_theta, rfr, skey, tau_m, v_rest, v_reset, ############################################################################ return _v, s, raw_s, _rfr -@partial(jit, static_argnums=[3, 4]) +#@partial(jit, static_argnums=[3, 4]) def _update_theta(dt, v_theta, s, tau_theta, theta_plus=0.05): ### Runs homeostatic threshold update dynamics one step (via Euler integration). #theta_decay = 0.9999999 #0.999999762 #jnp.exp(-dt/1e7) From f4c03a109d2ea69d356527a4222d3bd5cf5e68ad Mon Sep 17 00:00:00 2001 From: Alexander Ororbia Date: Sun, 28 Jul 2024 17:41:37 -0400 Subject: [PATCH 15/61] minor edit to dim-reduce --- ngclearn/utils/viz/dim_reduce.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/ngclearn/utils/viz/dim_reduce.py b/ngclearn/utils/viz/dim_reduce.py index 6a1c22e7..98646084 100755 --- a/ngclearn/utils/viz/dim_reduce.py +++ b/ngclearn/utils/viz/dim_reduce.py @@ -79,6 +79,8 @@ def plot_latents(code_vectors, labels, plot_fname="2Dcode_plot.jpg", alpha=1.): classes. plot_fname: /path/to/plot_fname. for saving the plot to disk + + alpha: """ curr_backend = plt.rcParams["backend"] matplotlib.use('Agg') ## temporarily go in Agg plt backend for tsne plotting @@ -91,7 +93,9 @@ def plot_latents(code_vectors, labels, plot_fname="2Dcode_plot.jpg", alpha=1.): lab = np.argmax(lab, 1) plt.figure(figsize=(8, 6)) plt.scatter(code_vectors[:, 0], code_vectors[:, 1], c=lab, cmap=cmap, alpha=alpha) - plt.colorbar() + colorbar = plt.colorbar() + #colorbar.set_alpha(1) + #plt.draw_all() plt.grid() plt.savefig("{0}".format(plot_fname), dpi=300) plt.clf() From 8d5bbd19037529e5a7baa5ccb7013d552db8c70e Mon Sep 17 00:00:00 2001 From: Will Gebhardt Date: Tue, 6 Aug 2024 14:25:27 -0400 Subject: [PATCH 16/61] updated monitor plot code --- ngclearn/components/base_monitor.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ngclearn/components/base_monitor.py b/ngclearn/components/base_monitor.py index 00c338b2..8238264d 100644 --- a/ngclearn/components/base_monitor.py +++ b/ngclearn/components/base_monitor.py @@ -310,4 +310,4 @@ def make_plot(self, compartment, ax=None, ylabel=None, xlabel=None, title=None, for k in range(n): _ax.plot(vals[:, 0, k]) else: - plot_func(vals, ax=_ax) + plot_func(vals[:, :, 0:n], ax=_ax) From 97c4d92d0791593fe4059824b03f8c3e8326c1a0 Mon Sep 17 00:00:00 2001 From: ago109 Date: Tue, 6 Aug 2024 14:29:23 -0400 Subject: [PATCH 17/61] update to dim-reduce --- ngclearn/utils/viz/dim_reduce.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/ngclearn/utils/viz/dim_reduce.py b/ngclearn/utils/viz/dim_reduce.py index 6a1c22e7..b37f5db6 100755 --- a/ngclearn/utils/viz/dim_reduce.py +++ b/ngclearn/utils/viz/dim_reduce.py @@ -79,6 +79,8 @@ def plot_latents(code_vectors, labels, plot_fname="2Dcode_plot.jpg", alpha=1.): classes. plot_fname: /path/to/plot_fname. for saving the plot to disk + + alpha: """ curr_backend = plt.rcParams["backend"] matplotlib.use('Agg') ## temporarily go in Agg plt backend for tsne plotting @@ -91,7 +93,9 @@ def plot_latents(code_vectors, labels, plot_fname="2Dcode_plot.jpg", alpha=1.): lab = np.argmax(lab, 1) plt.figure(figsize=(8, 6)) plt.scatter(code_vectors[:, 0], code_vectors[:, 1], c=lab, cmap=cmap, alpha=alpha) - plt.colorbar() + colorbar = plt.colorbar() + colorbar.set_alpha(1) + #plt.draw_all() plt.grid() plt.savefig("{0}".format(plot_fname), dpi=300) plt.clf() From 77f347fabbd755f66120f080728f2697186a6f82 Mon Sep 17 00:00:00 2001 From: ago109 Date: Tue, 6 Aug 2024 14:54:58 -0400 Subject: [PATCH 18/61] integrated phasor-cell, minor cleanup of latency --- ngclearn/components/__init__.py | 1 + .../components/input_encoders/__init__.py | 1 + .../components/input_encoders/latencyCell.py | 4 +- .../components/input_encoders/phasorCell.py | 179 ++++++++++++++++++ 4 files changed, 183 insertions(+), 2 deletions(-) create mode 100755 ngclearn/components/input_encoders/phasorCell.py diff --git a/ngclearn/components/__init__.py b/ngclearn/components/__init__.py index 005fbacc..1c35cfb1 100644 --- a/ngclearn/components/__init__.py +++ b/ngclearn/components/__init__.py @@ -21,6 +21,7 @@ from .input_encoders.bernoulliCell import BernoulliCell from .input_encoders.poissonCell import PoissonCell from .input_encoders.latencyCell import LatencyCell +from .input_encoders.phasorCell import PhasorCell ## point to synapse component types from .synapses.denseSynapse import DenseSynapse from .synapses.staticSynapse import StaticSynapse diff --git a/ngclearn/components/input_encoders/__init__.py b/ngclearn/components/input_encoders/__init__.py index 1be55d58..b779226e 100644 --- a/ngclearn/components/input_encoders/__init__.py +++ b/ngclearn/components/input_encoders/__init__.py @@ -1,3 +1,4 @@ from .bernoulliCell import BernoulliCell from .poissonCell import PoissonCell from .latencyCell import LatencyCell +from .phasorCell import PhasorCell diff --git a/ngclearn/components/input_encoders/latencyCell.py b/ngclearn/components/input_encoders/latencyCell.py index cf6db5ca..104dac37 100755 --- a/ngclearn/components/input_encoders/latencyCell.py +++ b/ngclearn/components/input_encoders/latencyCell.py @@ -48,7 +48,7 @@ def _calc_spike_times_linear(data, tau, thr, first_spk_t, num_steps=1., projected spike times """ _tau = tau - if normalize == True: + if normalize: _tau = num_steps - 1. - first_spk_t ## linear normalization #torch.clamp_max((-tau * (data - 1)), -tau * (threshold - 1)) stimes = -_tau * (data - 1.) ## calc raw latency code values @@ -85,7 +85,7 @@ def _calc_spike_times_nonlinear(data, tau, thr, first_spk_t, eps=1e-7, stimes = jnp.log(_data / (_data - thr)) * tau ## calc spike times stimes = stimes + first_spk_t - if normalize == True: + if normalize: term1 = (stimes - first_spk_t) term2 = (num_steps - first_spk_t - 1.) term3 = jnp.max(stimes - first_spk_t) diff --git a/ngclearn/components/input_encoders/phasorCell.py b/ngclearn/components/input_encoders/phasorCell.py new file mode 100755 index 00000000..68ea1517 --- /dev/null +++ b/ngclearn/components/input_encoders/phasorCell.py @@ -0,0 +1,179 @@ +from ngclearn import resolver, Compartment +from ngclearn.components.jaxComponent import JaxComponent +from ngclearn.utils import tensorstats +from jax import numpy as jnp, random +from ngcsimlib.logger import warn + +class PhasorCell(JaxComponent): + """ + A phasor cell that emits a pulse at a regular interval. + + | --- Cell Input Compartments: --- + | inputs - input (takes in external signals) + | --- Cell State Compartments: --- + | key - JAX PRNG key + | --- Cell Output Compartments: --- + | outputs - output + | tols - time-of-last-spike + + Args: + name: the string name of this cell + + n_units: number of cellular entities (neural population size) + + target_freq: maximum frequency (in Hertz) of this spike train + (must be > 0.) + """ + + # Define Functions + def __init__(self, name, n_units, target_freq=63.75, batch_size=1, + **kwargs): + super().__init__(name, **kwargs) + + ## Phasor meta-parameters + self.target_freq = target_freq ## maximum frequency (in Hertz/Hz) + + ## Layer Size Setup + self.batch_size = batch_size + self.n_units = n_units + _key, subkey = random.split(self.key.value, 2) + self.key.set(_key) + ## Compartment setup + restVals = jnp.zeros((self.batch_size, self.n_units)) + self.inputs = Compartment(restVals, + display_name="Input Stimulus") # input + # compartment + self.outputs = Compartment(restVals, + display_name="Spikes") # output compartment + self.tols = Compartment(initial_value=restVals, + display_name="Time-of-Last-Spike", units="ms") # time of last spike + self.angles = Compartment(restVals, display_name="Angles", units="deg") + # self.base_scale = random.uniform(subkey, self.angles.value.shape, + # minval=0.75, maxval=1.25) + # self.base_scale = ((random.normal(subkey, self.angles.value.shape) * 0.15) + 1) + # alpha = ((random.normal(subkey, self.angles.value.shape) * (jnp.sqrt(target_freq) / target_freq)) + 1) + # beta = random.poisson(subkey, lam=target_freq, shape=self.angles.value.shape) / target_freq + + self.base_scale = random.poisson(subkey, lam=target_freq, shape=self.angles.value.shape) / target_freq + + def validate(self, dt=None, **validation_kwargs): + valid = super().validate(**validation_kwargs) + if dt is None: + warn(f"{self.name} requires a validation kwarg of `dt`") + return False + ## check for unstable combinations of dt and target-frequency + # meta-params + events_per_timestep = ( + dt / 1000.) * self.target_freq ## + # compute scaled probability + if events_per_timestep > 1.: + valid = False + warn( + f"{self.name} will be unable to make as many temporal events " + f"as " + f"requested! ({events_per_timestep} events/timestep) Unstable " + f"combination of dt = {dt} and target_freq = " + f"{self.target_freq} " + f"being used!" + ) + return valid + + @staticmethod + def _advance_state(t, dt, target_freq, key, + inputs, angles, tols, base_scale): + ms_per_second = 1000 # ms/s + events_per_ms = target_freq / ms_per_second # e/s s/ms -> e/ms + ms_per_event = 1 / events_per_ms # ms/e + time_step_per_event = ms_per_event / dt # ms/e * ts/ms -> ts / e + angle_per_event = 2 * jnp.pi # rad / e + angle_per_timestep = angle_per_event / time_step_per_event # rad / e + # * e/ts -> rad / ts + key, subkey = random.split(key, 2) + # scatter = random.uniform(subkey, angles.shape, minval=0.5, + # maxval=1.5) * base_scale + + scatter = ((random.normal(subkey, angles.shape) * 0.2) + 1) * base_scale + scattered_update = angle_per_timestep * scatter + scaled_scattered_update = scattered_update * inputs + + updated_angles = angles + scaled_scattered_update + outputs = jnp.where(updated_angles > angle_per_event, 1., 0.) + updated_angles = jnp.where(updated_angles > angle_per_event, + updated_angles - angle_per_event, + updated_angles) + tols = tols * (1. - outputs) + t * outputs + + return outputs, tols, key, updated_angles + + @resolver(_advance_state) + def advance_state(self, outputs, tols, key, angles): + self.outputs.set(outputs) + self.tols.set(tols) + self.key.set(key) + self.angles.set(angles) + + @staticmethod + def _reset(batch_size, n_units, key, target_freq): + restVals = jnp.zeros((batch_size, n_units)) + key, subkey = random.split(key, 2) + return restVals, restVals, restVals, restVals, key + + @resolver(_reset) + def reset(self, inputs, outputs, tols, angles, key): + self.inputs.set(inputs) + self.outputs.set(outputs) + self.tols.set(tols) + self.key.set(key) + self.angles.set(angles) + + def save(self, directory, **kwargs): + file_name = directory + "/" + self.name + ".npz" + jnp.savez(file_name, key=self.key.value) + + def load(self, directory, **kwargs): + file_name = directory + "/" + self.name + ".npz" + data = jnp.load(file_name) + self.key.set(data['key']) + + @classmethod + def help(cls): ## component help function + properties = { + "cell_type": "Phasor - Produces input at a fairly regular " + "intervals with small amounts of noise)" + } + compartment_props = { + "inputs": + {"inputs": "Takes in external input signal values"}, + "states": + {"key": "JAX PRNG key", + "angles": "The current angle of the phasor"}, + "outputs": + {"tols": "Time-of-last-spike", + "outputs": "Binary spike values emitted at time t"}, + } + hyperparams = { + "n_units": "Number of neuronal cells to model in this layer", + "batch_size": "Batch size dimension of this component", + "target_freq": "Maximum spike frequency of the train produced", + } + info = {cls.__name__: properties, + "compartments": compartment_props, + "hyperparameters": hyperparams} + return info + + def __repr__(self): + comps = [varname for varname in dir(self) if + Compartment.is_compartment(getattr(self, varname))] + maxlen = max(len(c) for c in comps) + 5 + lines = f"[{self.__class__.__name__}] PATH: {self.name}\n" + for c in comps: + stats = tensorstats(getattr(self, c).value) + if stats is not None: + line = [f"{k}: {v}" for k, v in stats.items()] + line = ", ".join(line) + else: + line = "None" + lines += f" {f'({c})'.ljust(maxlen)}{line}\n" + return lines + + From 714a58c509b9a3d68a2a0ecfcc031f7e661f0469 Mon Sep 17 00:00:00 2001 From: ago109 Date: Wed, 7 Aug 2024 16:21:54 -0400 Subject: [PATCH 19/61] tweak to adex thr arg --- ngclearn/components/neurons/spiking/adExCell.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/ngclearn/components/neurons/spiking/adExCell.py b/ngclearn/components/neurons/spiking/adExCell.py index ef0c9450..7e3b0547 100755 --- a/ngclearn/components/neurons/spiking/adExCell.py +++ b/ngclearn/components/neurons/spiking/adExCell.py @@ -112,7 +112,7 @@ class AdExCell(JaxComponent): intrinsic_mem_thr: intrinsic membrane threshold (Default: -55 mV) - v_thr: voltage/membrane threshold (to obtain action potentials in terms + thr: voltage/membrane threshold (to obtain action potentials in terms of binary spikes) (Default: 5 mV) v_rest: membrane resting potential (Default: -72 mV) @@ -136,7 +136,7 @@ class AdExCell(JaxComponent): # Define Functions def __init__(self, name, n_units, tau_m=15., resist_m=1., tau_w=400., - v_sharpness=2., intrinsic_mem_thr=-55., v_thr=5., v_rest=-72., + v_sharpness=2., intrinsic_mem_thr=-55., thr=5., v_rest=-72., v_reset=-75., a=0.1, b=0.75, v0=-70., w0=0., integration_type="euler", batch_size=1, **kwargs): super().__init__(name, **kwargs) @@ -158,7 +158,7 @@ def __init__(self, name, n_units, tau_m=15., resist_m=1., tau_w=400., self.v0 = v0 ## initial membrane potential/voltage condition self.w0 = w0 ## initial w-parameter condition - self.v_thr = v_thr + self.thr = thr ## Layer Size Setup self.batch_size = batch_size @@ -174,9 +174,9 @@ def __init__(self, name, n_units, tau_m=15., resist_m=1., tau_w=400., units="ms") ## time-of-last-spike @staticmethod - def _advance_state(t, dt, tau_m, R_m, tau_w, v_thr, a, b, sharpV, vT, + def _advance_state(t, dt, tau_m, R_m, tau_w, thr, a, b, sharpV, vT, v_rest, v_reset, intgFlag, j, v, w, tols): - v, w, s = _run_cell(dt, j, v, w, v_thr, tau_m, tau_w, a, b, sharpV, vT, + v, w, s = _run_cell(dt, j, v, w, thr, tau_m, tau_w, a, b, sharpV, vT, v_rest, v_reset, R_m, intgFlag) tols = _update_times(t, s, tols) return j, v, w, s, tols @@ -230,7 +230,7 @@ def help(cls): ## component help function "tau_m": "Cell membrane time constant", "resist_m": "Membrane resistance value", "tau_w": "Recovery variable time constant", - "v_thr": "Base voltage threshold value", + "thr": "Base voltage threshold value", "v_rest": "Resting membrane potential value", "v_reset": "Reset membrane potential value", "v_sharpness": "Slope factor/voltage sharpness constant", From 6ec2e7a055face85b5731cd2c3f03bd4f7227f9e Mon Sep 17 00:00:00 2001 From: ago109 Date: Wed, 7 Aug 2024 16:24:32 -0400 Subject: [PATCH 20/61] tweak to adex thr arg --- ngclearn/components/neurons/spiking/adExCell.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/ngclearn/components/neurons/spiking/adExCell.py b/ngclearn/components/neurons/spiking/adExCell.py index 7e3b0547..699061a2 100755 --- a/ngclearn/components/neurons/spiking/adExCell.py +++ b/ngclearn/components/neurons/spiking/adExCell.py @@ -2,6 +2,7 @@ from ngclearn import resolver, Component, Compartment from ngclearn.components.jaxComponent import JaxComponent from ngclearn.utils import tensorstats +from ngcsimlib.deprecators import deprecate_args from ngclearn.utils.diffeq.ode_utils import get_integrator_code, \ step_euler, step_rk2 @@ -134,7 +135,7 @@ class AdExCell(JaxComponent): at an increase in computational cost (and simulation time) """ - # Define Functions + @deprecate_args(v_thr="thr") def __init__(self, name, n_units, tau_m=15., resist_m=1., tau_w=400., v_sharpness=2., intrinsic_mem_thr=-55., thr=5., v_rest=-72., v_reset=-75., a=0.1, b=0.75, v0=-70., w0=0., From fb8524ae080cbceab5b28e696891e6d2466ae223 Mon Sep 17 00:00:00 2001 From: ago109 Date: Wed, 7 Aug 2024 22:41:29 -0400 Subject: [PATCH 21/61] integrated resonate-and-fire neuronal cell --- ngclearn/components/__init__.py | 1 + ngclearn/components/neurons/__init__.py | 1 + .../components/neurons/spiking/RAFCell.py | 251 ++++++++++++++++++ .../components/neurons/spiking/__init__.py | 1 + 4 files changed, 254 insertions(+) create mode 100755 ngclearn/components/neurons/spiking/RAFCell.py diff --git a/ngclearn/components/__init__.py b/ngclearn/components/__init__.py index 1c35cfb1..88808f78 100644 --- a/ngclearn/components/__init__.py +++ b/ngclearn/components/__init__.py @@ -13,6 +13,7 @@ from .neurons.spiking.adExCell import AdExCell from .neurons.spiking.fitzhughNagumoCell import FitzhughNagumoCell from .neurons.spiking.izhikevichCell import IzhikevichCell +from .neurons.spiking.RAFCell import RAFCell ## point to transformer/operater component types from .other.varTrace import VarTrace from .other.expKernel import ExpKernel diff --git a/ngclearn/components/neurons/__init__.py b/ngclearn/components/neurons/__init__.py index 42a4a971..6312b379 100644 --- a/ngclearn/components/neurons/__init__.py +++ b/ngclearn/components/neurons/__init__.py @@ -12,3 +12,4 @@ from .spiking.adExCell import AdExCell from .spiking.fitzhughNagumoCell import FitzhughNagumoCell from .spiking.izhikevichCell import IzhikevichCell +from .spiking.RAFCell import RAFCell diff --git a/ngclearn/components/neurons/spiking/RAFCell.py b/ngclearn/components/neurons/spiking/RAFCell.py new file mode 100755 index 00000000..6758daf5 --- /dev/null +++ b/ngclearn/components/neurons/spiking/RAFCell.py @@ -0,0 +1,251 @@ +from jax import numpy as jnp, jit +from ngclearn import resolver, Component, Compartment +from ngclearn.components.jaxComponent import JaxComponent +from ngclearn.utils import tensorstats +from ngclearn.utils.diffeq.ode_utils import get_integrator_code, \ + step_euler, step_rk2 + +@jit +def _update_times(t, s, tols): + """ + Updates time-of-last-spike (tols) variable. + + Args: + t: current time (a scalar/int value) + + s: binary spike vector + + tols: current time-of-last-spike variable + + Returns: + updated tols variable + """ + _tols = (1. - s) * tols + (s * t) + return _tols + +@jit +def _dfv_internal(j, v, w, tau_m, omega, b): ## "voltage" dynamics + # dy/dt = omega x + b y + dv_dt = omega * w + v * b ## dv/dt + dv_dt = dv_dt * (1./tau_m) + return dv_dt + +def _dfv(t, v, params): ## voltage dynamics wrapper + j, w, tau_m, omega, b = params + dv_dt = _dfv_internal(j, v, w, tau_m, omega, b) + return dv_dt + +@jit +def _dfw_internal(j, v, w, tau_w, omega, b): ## raw angular driver dynamics + # dx/dt = b x − omega y + I + dw_dt = w * b - v * omega + j + dw_dt = dw_dt * (1./tau_w) + return dw_dt + +def _dfw(t, w, params): ## angular driver dynamics wrapper + j, v, tau_w, omega, b = params + dv_dt = _dfw_internal(j, v, w, tau_w, omega, b) + return dv_dt + +@jit +def _emit_spike(v, v_thr): + s = (v > v_thr).astype(jnp.float32) + return s + +class RAFCell(JaxComponent): + """ + The resonate-and-fire (RAF) neuronal cell + model; a two-variable model. This cell model iteratively evolves + voltage "v" and angular driver "w". + + The specific pair of differential equations that characterize this cell + are (for adjusting v and w, given current j, over time): + + | tau_m * dv/dt = -(v - v_rest) + sharpV * exp((v - vT)/sharpV) - R_m * w + R_m * j + | tau_w * dw/dt = -w + (v - v_rest) * a + | where w = w + s * (w + b) [in the event of a spike] + + | --- Cell Input Compartments: --- + | j - electrical current input (takes in external signals) + | --- Cell State Compartments: --- + | v - membrane potential/voltage state + | w - angular driver variable state + | key - JAX PRNG key + | --- Cell Output Compartments: --- + | s - emitted binary spikes/action potentials + | tols - time-of-last-spike + + | References: + | Izhikevich, Eugene M. "Resonate-and-fire neurons." Neural networks + | 14.6-7 (2001): 883-894. + + Args: + name: the string name of this cell + + n_units: number of cellular entities (neural population size) + + tau_m: membrane time constant (Default: 15 ms) + + resist_m: membrane resistance (Default: 1 mega-Ohm) + + tau_w: angular driver variable time constant (Default: 400 ms) + + thr: voltage/membrane threshold (to obtain action potentials in terms + of binary spikes) (Default: 5 mV) + + v_rest: membrane resting potential (Default: -72 mV) + + b: oscillation dampening factor (Default: -1.) + + v0: initial condition / reset for voltage (Default: -70 mV) + + w0: initial condition / reset for angular driver (Default: 0 mV) + + integration_type: type of integration to use for this cell's dynamics; + current supported forms include "euler" (Euler/RK-1 integration) + and "midpoint" or "rk2" (midpoint method/RK-2 integration) (Default: "euler") + + :Note: setting the integration type to the midpoint method will + increase the accuray of the estimate of the cell's evolution + at an increase in computational cost (and simulation time) + """ + + # Define Functions + def __init__(self, name, n_units, tau_m=15., resist_m=1., tau_w=400., + omega=10., thr=5., v_rest=-72., + v_reset=-75., w_reset=0., b=-1., v0=-70., w0=0., + integration_type="euler", batch_size=1, **kwargs): + super().__init__(name, **kwargs) + + ## Integration properties + self.integrationType = integration_type + self.intgFlag = get_integrator_code(self.integrationType) + + ## Cell properties + self.tau_m = tau_m + self.R_m = resist_m + self.tau_w = tau_w + self.omega = omega ## angular frequency + self.b = b ## dampening factor + ## note: the smaller b is, the faster the oscillation dampens to resting state values + self.v_rest = v_rest + self.v_reset = v_reset + self.w_reset = w_reset + + self.v0 = v0 ## initial membrane potential/voltage condition + self.w0 = w0 ## initial w-parameter condition + self.thr = thr + + ## Layer Size Setup + self.batch_size = batch_size + self.n_units = n_units + + ## Compartment setup + restVals = jnp.zeros((self.batch_size, self.n_units)) + self.j = Compartment(restVals, display_name="Current", units="mA") + self.v = Compartment(restVals + self.v0, display_name="Voltage", units="mV") + self.w = Compartment(restVals + self.w0, display_name="Angular-Driver") + self.s = Compartment(restVals, display_name="Spikes") + self.tols = Compartment(restVals, display_name="Time-of-Last-Spike", + units="ms") ## time-of-last-spike + + @staticmethod + def _advance_state(t, dt, tau_m, R_m, tau_w, thr, omega, b, v_rest, + v_reset, w_reset, intgFlag, j, v, w, tols): + j_ = j * R_m + if intgFlag == 1: ## RK-2/midpoint + w_params = (j_, v, tau_w, omega, b) + _, _w = step_rk2(0., w, _dfw, dt, w_params) + v_params = (j_, w, tau_m, omega, b) + _, _v = step_rk2(0., v, _dfv, dt, v_params) + else: # integType == 0 (default -- Euler) + w_params = (j_, v, tau_w, omega, b) + _, _w = step_euler(0., w, _dfw, dt, w_params) + v_params = (j_, w, tau_m, omega, b) + _, _v = step_euler(0., v, _dfv, dt, v_params) + s = _emit_spike(_v, thr) + ## hyperpolarize/reset/snap variables + v = _v * (1. - s) + s * v_reset + w = _w * (1. - s) + s * w_reset + + tols = _update_times(t, s, tols) + return j, v, w, s, tols + + @resolver(_advance_state) + def advance_state(self, j, v, w, s, tols): + self.j.set(j) + self.w.set(w) + self.v.set(v) + self.s.set(s) + self.tols.set(tols) + + @staticmethod + def _reset(batch_size, n_units, v0, w0): + restVals = jnp.zeros((batch_size, n_units)) + j = restVals # None + v = restVals + v0 + w = restVals + w0 + s = restVals #+ 0 + tols = restVals #+ 0 + return j, v, w, s, tols + + @resolver(_reset) + def reset(self, j, v, w, s, tols): + self.j.set(j) + self.v.set(v) + self.w.set(w) + self.s.set(s) + self.tols.set(tols) + + @classmethod + def help(cls): ## component help function + properties = { + "cell_type": "RAFCell - evolves neurons according to nonlinear, " + "resonate-and-fire dual-ODE spiking cell dynamics." + } + compartment_props = { + "inputs": + {"j": "External input electrical current", + "key": "JAX PRNG key"}, + "states": + {"v": "Membrane potential/voltage at time t", + "w": "Recovery variable at time t"}, + "outputs": + {"s": "Emitted spikes/pulses at time t", + "tols": "Time-of-last-spike"}, + } + hyperparams = { + "n_units": "Number of neuronal cells to model in this layer", + "batch_size": "Batch size dimension of this component", + "tau_m": "Cell membrane time constant", + "resist_m": "Membrane resistance value", + "tau_w": "Recovery variable time constant", + "v_thr": "Base voltage threshold value", + "v_rest": "Resting membrane potential value", + "v_reset": "Reset membrane potential value", + "b": "Exponential dampening factor applied to oscillations", + "omega": "Angular frequency of neuronal progress per second (radians)", + "v0": "Initial condition for membrane potential/voltage", + "w0": "Initial condition for membrane angular driver variable", + "integration_type": "Type of numerical integration to use for the cell dynamics" + } + info = {cls.__name__: properties, + "compartments": compartment_props, + "dynamics": "tau_m * dv/dt = omega * w + v * b; " + "tau_w * dw/dt = w * b - v * omega + j", + "hyperparameters": hyperparams} + return info + + def __repr__(self): + comps = [varname for varname in dir(self) if Compartment.is_compartment(getattr(self, varname))] + maxlen = max(len(c) for c in comps) + 5 + lines = f"[{self.__class__.__name__}] PATH: {self.name}\n" + for c in comps: + stats = tensorstats(getattr(self, c).value) + if stats is not None: + line = [f"{k}: {v}" for k, v in stats.items()] + line = ", ".join(line) + else: + line = "None" + lines += f" {f'({c})'.ljust(maxlen)}{line}\n" + return lines diff --git a/ngclearn/components/neurons/spiking/__init__.py b/ngclearn/components/neurons/spiking/__init__.py index 2934eda9..766d44b5 100644 --- a/ngclearn/components/neurons/spiking/__init__.py +++ b/ngclearn/components/neurons/spiking/__init__.py @@ -7,3 +7,4 @@ from .adExCell import AdExCell from .fitzhughNagumoCell import FitzhughNagumoCell from .izhikevichCell import IzhikevichCell +from .RAFCell import RAFCell From dd49e5fe81abaf86460a2ee80b1c3d2a78040d7f Mon Sep 17 00:00:00 2001 From: ago109 Date: Wed, 7 Aug 2024 22:58:11 -0400 Subject: [PATCH 22/61] mod to raf-cell --- ngclearn/components/neurons/spiking/RAFCell.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ngclearn/components/neurons/spiking/RAFCell.py b/ngclearn/components/neurons/spiking/RAFCell.py index 6758daf5..a1a03cd8 100755 --- a/ngclearn/components/neurons/spiking/RAFCell.py +++ b/ngclearn/components/neurons/spiking/RAFCell.py @@ -123,7 +123,7 @@ def __init__(self, name, n_units, tau_m=15., resist_m=1., tau_w=400., ## Cell properties self.tau_m = tau_m - self.R_m = resist_m + self.resist_m = resist_m self.tau_w = tau_w self.omega = omega ## angular frequency self.b = b ## dampening factor @@ -150,9 +150,9 @@ def __init__(self, name, n_units, tau_m=15., resist_m=1., tau_w=400., units="ms") ## time-of-last-spike @staticmethod - def _advance_state(t, dt, tau_m, R_m, tau_w, thr, omega, b, v_rest, + def _advance_state(t, dt, tau_m, resist_m, tau_w, thr, omega, b, v_rest, v_reset, w_reset, intgFlag, j, v, w, tols): - j_ = j * R_m + j_ = j * resist_m if intgFlag == 1: ## RK-2/midpoint w_params = (j_, v, tau_w, omega, b) _, _w = step_rk2(0., w, _dfw, dt, w_params) From 888220804af2cbefab017bf1fddf0f84b909bf20 Mon Sep 17 00:00:00 2001 From: ago109 Date: Thu, 8 Aug 2024 17:46:18 -0400 Subject: [PATCH 23/61] cleaned up raf --- .../components/neurons/spiking/RAFCell.py | 43 ++++++++++--------- 1 file changed, 22 insertions(+), 21 deletions(-) diff --git a/ngclearn/components/neurons/spiking/RAFCell.py b/ngclearn/components/neurons/spiking/RAFCell.py index a1a03cd8..7c1b946b 100755 --- a/ngclearn/components/neurons/spiking/RAFCell.py +++ b/ngclearn/components/neurons/spiking/RAFCell.py @@ -61,9 +61,9 @@ class RAFCell(JaxComponent): The specific pair of differential equations that characterize this cell are (for adjusting v and w, given current j, over time): - | tau_m * dv/dt = -(v - v_rest) + sharpV * exp((v - vT)/sharpV) - R_m * w + R_m * j - | tau_w * dw/dt = -w + (v - v_rest) * a - | where w = w + s * (w + b) [in the event of a spike] + | tau_m * dv/dt = omega * w + v * b + | tau_w * dw/dt = w * b - v * omega + j + | where omega is angular frequency (Hz) and b is exponential dampening factor | --- Cell Input Compartments: --- | j - electrical current input (takes in external signals) @@ -93,13 +93,11 @@ class RAFCell(JaxComponent): thr: voltage/membrane threshold (to obtain action potentials in terms of binary spikes) (Default: 5 mV) - v_rest: membrane resting potential (Default: -72 mV) + v_reset: membrane reset potential condition (Default: 0 mV) - b: oscillation dampening factor (Default: -1.) - - v0: initial condition / reset for voltage (Default: -70 mV) + w_reset: reset condition for angular driver (Default: 0 mV) - w0: initial condition / reset for angular driver (Default: 0 mV) + b: oscillation dampening factor (Default: -1.) integration_type: type of integration to use for this cell's dynamics; current supported forms include "euler" (Euler/RK-1 integration) @@ -112,9 +110,9 @@ class RAFCell(JaxComponent): # Define Functions def __init__(self, name, n_units, tau_m=15., resist_m=1., tau_w=400., - omega=10., thr=5., v_rest=-72., - v_reset=-75., w_reset=0., b=-1., v0=-70., w0=0., + omega=10., thr=5., v_reset=0., w_reset=0., b=-1., integration_type="euler", batch_size=1, **kwargs): + #v_rest=-72., v_reset=-75., w_reset=0., thr=5., v0=-70., w0=0., super().__init__(name, **kwargs) ## Integration properties @@ -128,12 +126,9 @@ def __init__(self, name, n_units, tau_m=15., resist_m=1., tau_w=400., self.omega = omega ## angular frequency self.b = b ## dampening factor ## note: the smaller b is, the faster the oscillation dampens to resting state values - self.v_rest = v_rest + #self.v_rest = v_rest self.v_reset = v_reset self.w_reset = w_reset - - self.v0 = v0 ## initial membrane potential/voltage condition - self.w0 = w0 ## initial w-parameter condition self.thr = thr ## Layer Size Setup @@ -150,8 +145,12 @@ def __init__(self, name, n_units, tau_m=15., resist_m=1., tau_w=400., units="ms") ## time-of-last-spike @staticmethod - def _advance_state(t, dt, tau_m, resist_m, tau_w, thr, omega, b, v_rest, + def _advance_state(t, dt, tau_m, resist_m, tau_w, thr, omega, b, v_reset, w_reset, intgFlag, j, v, w, tols): + ## center variables before running dynamics + v = v - v_reset + w = w - w_reset + ## continue with centered dynamics j_ = j * resist_m if intgFlag == 1: ## RK-2/midpoint w_params = (j_, v, tau_w, omega, b) @@ -165,9 +164,11 @@ def _advance_state(t, dt, tau_m, resist_m, tau_w, thr, omega, b, v_rest, _, _v = step_euler(0., v, _dfv, dt, v_params) s = _emit_spike(_v, thr) ## hyperpolarize/reset/snap variables - v = _v * (1. - s) + s * v_reset - w = _w * (1. - s) + s * w_reset - + v = _v * (1. - s) + s #* v_reset + w = _w * (1. - s) + s #* w_reset + ## artificially shift variables back to rest/reset values + v = v + v_reset + w = w + w_reset tols = _update_times(t, s, tols) return j, v, w, s, tols @@ -180,11 +181,11 @@ def advance_state(self, j, v, w, s, tols): self.tols.set(tols) @staticmethod - def _reset(batch_size, n_units, v0, w0): + def _reset(batch_size, n_units, v_reset, w_reset): restVals = jnp.zeros((batch_size, n_units)) j = restVals # None - v = restVals + v0 - w = restVals + w0 + v = restVals + v_reset + w = restVals + w_reset s = restVals #+ 0 tols = restVals #+ 0 return j, v, w, s, tols From ee50f333cd728ee51b12289b53e94a3074117fc6 Mon Sep 17 00:00:00 2001 From: ago109 Date: Thu, 8 Aug 2024 17:52:46 -0400 Subject: [PATCH 24/61] cleaned up raf --- .../components/neurons/spiking/RAFCell.py | 35 ++++++++++--------- 1 file changed, 18 insertions(+), 17 deletions(-) diff --git a/ngclearn/components/neurons/spiking/RAFCell.py b/ngclearn/components/neurons/spiking/RAFCell.py index 7c1b946b..0f23107c 100755 --- a/ngclearn/components/neurons/spiking/RAFCell.py +++ b/ngclearn/components/neurons/spiking/RAFCell.py @@ -93,11 +93,17 @@ class RAFCell(JaxComponent): thr: voltage/membrane threshold (to obtain action potentials in terms of binary spikes) (Default: 5 mV) + omega: angular frequency (Default: 10) + + b: oscillation dampening factor (Default: -1) + v_reset: membrane reset potential condition (Default: 0 mV) - w_reset: reset condition for angular driver (Default: 0 mV) + w_reset: reset condition for angular driver (Default: 0) + + v0: membrane potential initial condition (Default: 0 mV) - b: oscillation dampening factor (Default: -1.) + w0: angular driver initial condition (Default: 0) integration_type: type of integration to use for this cell's dynamics; current supported forms include "euler" (Euler/RK-1 integration) @@ -110,8 +116,8 @@ class RAFCell(JaxComponent): # Define Functions def __init__(self, name, n_units, tau_m=15., resist_m=1., tau_w=400., - omega=10., thr=5., v_reset=0., w_reset=0., b=-1., - integration_type="euler", batch_size=1, **kwargs): + thr=5., omega=10., b=-1., v_reset=0., w_reset=0., + v0=0., w0=0., integration_type="euler", batch_size=1, **kwargs): #v_rest=-72., v_reset=-75., w_reset=0., thr=5., v0=-70., w0=0., super().__init__(name, **kwargs) @@ -129,6 +135,8 @@ def __init__(self, name, n_units, tau_m=15., resist_m=1., tau_w=400., #self.v_rest = v_rest self.v_reset = v_reset self.w_reset = w_reset + self.v0 = v0 + self.w0 = w0 self.thr = thr ## Layer Size Setup @@ -147,9 +155,6 @@ def __init__(self, name, n_units, tau_m=15., resist_m=1., tau_w=400., @staticmethod def _advance_state(t, dt, tau_m, resist_m, tau_w, thr, omega, b, v_reset, w_reset, intgFlag, j, v, w, tols): - ## center variables before running dynamics - v = v - v_reset - w = w - w_reset ## continue with centered dynamics j_ = j * resist_m if intgFlag == 1: ## RK-2/midpoint @@ -164,11 +169,8 @@ def _advance_state(t, dt, tau_m, resist_m, tau_w, thr, omega, b, _, _v = step_euler(0., v, _dfv, dt, v_params) s = _emit_spike(_v, thr) ## hyperpolarize/reset/snap variables - v = _v * (1. - s) + s #* v_reset - w = _w * (1. - s) + s #* w_reset - ## artificially shift variables back to rest/reset values - v = v + v_reset - w = w + w_reset + v = _v * (1. - s) + s * v_reset + w = _w * (1. - s) + s * w_reset tols = _update_times(t, s, tols) return j, v, w, s, tols @@ -181,11 +183,11 @@ def advance_state(self, j, v, w, s, tols): self.tols.set(tols) @staticmethod - def _reset(batch_size, n_units, v_reset, w_reset): + def _reset(batch_size, n_units, v0, w0): restVals = jnp.zeros((batch_size, n_units)) j = restVals # None - v = restVals + v_reset - w = restVals + w_reset + v = restVals + v0 + w = restVals + w0 s = restVals #+ 0 tols = restVals #+ 0 return j, v, w, s, tols @@ -221,9 +223,8 @@ def help(cls): ## component help function "tau_m": "Cell membrane time constant", "resist_m": "Membrane resistance value", "tau_w": "Recovery variable time constant", - "v_thr": "Base voltage threshold value", - "v_rest": "Resting membrane potential value", "v_reset": "Reset membrane potential value", + "w_reset": "Reset angular driver value", "b": "Exponential dampening factor applied to oscillations", "omega": "Angular frequency of neuronal progress per second (radians)", "v0": "Initial condition for membrane potential/voltage", From 611e5b3c275731bf62273afd944e790062f11c31 Mon Sep 17 00:00:00 2001 From: ago109 Date: Fri, 9 Aug 2024 12:34:06 -0400 Subject: [PATCH 25/61] cleaned up raf-cell --- .../components/neurons/spiking/RAFCell.py | 62 ++++++++++--------- 1 file changed, 32 insertions(+), 30 deletions(-) diff --git a/ngclearn/components/neurons/spiking/RAFCell.py b/ngclearn/components/neurons/spiking/RAFCell.py index 0f23107c..af49f00e 100755 --- a/ngclearn/components/neurons/spiking/RAFCell.py +++ b/ngclearn/components/neurons/spiking/RAFCell.py @@ -2,6 +2,7 @@ from ngclearn import resolver, Component, Compartment from ngclearn.components.jaxComponent import JaxComponent from ngclearn.utils import tensorstats +from ngcsimlib.deprecators import deprecate_args from ngclearn.utils.diffeq.ode_utils import get_integrator_code, \ step_euler, step_rk2 @@ -37,7 +38,7 @@ def _dfv(t, v, params): ## voltage dynamics wrapper @jit def _dfw_internal(j, v, w, tau_w, omega, b): ## raw angular driver dynamics - # dx/dt = b x − omega y + I + # dx/dt = b x − omega y + I; I is scaled injected electrical current dw_dt = w * b - v * omega + j dw_dt = dw_dt * (1./tau_w) return dw_dt @@ -61,9 +62,10 @@ class RAFCell(JaxComponent): The specific pair of differential equations that characterize this cell are (for adjusting v and w, given current j, over time): - | tau_m * dv/dt = omega * w + v * b | tau_w * dw/dt = w * b - v * omega + j + | tau_v * dv/dt = omega * w + v * b | where omega is angular frequency (Hz) and b is exponential dampening factor + | Note: injected current j should generally be scaled by tau_w/dt | --- Cell Input Compartments: --- | j - electrical current input (takes in external signals) @@ -84,27 +86,27 @@ class RAFCell(JaxComponent): n_units: number of cellular entities (neural population size) - tau_m: membrane time constant (Default: 15 ms) + tau_v: membrane/voltage time constant (Default: 1 ms) - resist_m: membrane resistance (Default: 1 mega-Ohm) - - tau_w: angular driver variable time constant (Default: 400 ms) + tau_w: angular driver variable time constant (Default: 1 ms) thr: voltage/membrane threshold (to obtain action potentials in terms - of binary spikes) (Default: 5 mV) + of binary spikes) (Default: 1 mV) omega: angular frequency (Default: 10) b: oscillation dampening factor (Default: -1) - v_reset: membrane reset potential condition (Default: 0 mV) + v_reset: membrane potential reset condition (Default: 1 mV) - w_reset: reset condition for angular driver (Default: 0) + w_reset: reset condition for angular current driver (Default: 0) - v0: membrane potential initial condition (Default: 0 mV) + v0: membrane potential initial condition (Default: 1 mV) w0: angular driver initial condition (Default: 0) + resist_v: membrane resistance (Default: 1 mega-Ohm) + integration_type: type of integration to use for this cell's dynamics; current supported forms include "euler" (Euler/RK-1 integration) and "midpoint" or "rk2" (midpoint method/RK-2 integration) (Default: "euler") @@ -114,11 +116,11 @@ class RAFCell(JaxComponent): at an increase in computational cost (and simulation time) """ - # Define Functions - def __init__(self, name, n_units, tau_m=15., resist_m=1., tau_w=400., - thr=5., omega=10., b=-1., v_reset=0., w_reset=0., - v0=0., w0=0., integration_type="euler", batch_size=1, **kwargs): - #v_rest=-72., v_reset=-75., w_reset=0., thr=5., v0=-70., w0=0., + @deprecate_args(resist_m="resist_v", tau_m="tau_v") + def __init__(self, name, n_units, tau_v=1., tau_w=1., thr=1., omega=10., + b=-1., v_reset=1., w_reset=0., v0=0., w0=0., resist_v=1., + integration_type="euler", batch_size=1, **kwargs): + #v_rest=-72., v_reset=-75., w_reset=0., thr=5., v0=-70., w0=0., tau_w=400., thr=5., omega=10., b=-1. super().__init__(name, **kwargs) ## Integration properties @@ -126,13 +128,12 @@ def __init__(self, name, n_units, tau_m=15., resist_m=1., tau_w=400., self.intgFlag = get_integrator_code(self.integrationType) ## Cell properties - self.tau_m = tau_m - self.resist_m = resist_m + self.tau_v = tau_v + self.resist_v = resist_v self.tau_w = tau_w self.omega = omega ## angular frequency self.b = b ## dampening factor ## note: the smaller b is, the faster the oscillation dampens to resting state values - #self.v_rest = v_rest self.v_reset = v_reset self.w_reset = w_reset self.v0 = v0 @@ -153,24 +154,25 @@ def __init__(self, name, n_units, tau_m=15., resist_m=1., tau_w=400., units="ms") ## time-of-last-spike @staticmethod - def _advance_state(t, dt, tau_m, resist_m, tau_w, thr, omega, b, + def _advance_state(t, dt, tau_v, resist_v, tau_w, thr, omega, b, v_reset, w_reset, intgFlag, j, v, w, tols): ## continue with centered dynamics - j_ = j * resist_m + j_ = j * resist_v if intgFlag == 1: ## RK-2/midpoint w_params = (j_, v, tau_w, omega, b) _, _w = step_rk2(0., w, _dfw, dt, w_params) - v_params = (j_, w, tau_m, omega, b) + v_params = (j_, w, tau_v, omega, b) _, _v = step_rk2(0., v, _dfv, dt, v_params) else: # integType == 0 (default -- Euler) w_params = (j_, v, tau_w, omega, b) _, _w = step_euler(0., w, _dfw, dt, w_params) - v_params = (j_, w, tau_m, omega, b) + v_params = (j_, w, tau_v, omega, b) _, _v = step_euler(0., v, _dfv, dt, v_params) s = _emit_spike(_v, thr) ## hyperpolarize/reset/snap variables - v = _v * (1. - s) + s * v_reset w = _w * (1. - s) + s * w_reset + v = _v * (1. - s) + s * v_reset + tols = _update_times(t, s, tols) return j, v, w, s, tols @@ -183,11 +185,11 @@ def advance_state(self, j, v, w, s, tols): self.tols.set(tols) @staticmethod - def _reset(batch_size, n_units, v0, w0): + def _reset(batch_size, n_units, v_reset, w_reset): restVals = jnp.zeros((batch_size, n_units)) j = restVals # None - v = restVals + v0 - w = restVals + w0 + v = restVals + v_reset + w = restVals + w_reset s = restVals #+ 0 tols = restVals #+ 0 return j, v, w, s, tols @@ -212,7 +214,7 @@ def help(cls): ## component help function "key": "JAX PRNG key"}, "states": {"v": "Membrane potential/voltage at time t", - "w": "Recovery variable at time t"}, + "w": "Angular current driver variable at time t"}, "outputs": {"s": "Emitted spikes/pulses at time t", "tols": "Time-of-last-spike"}, @@ -220,8 +222,7 @@ def help(cls): ## component help function hyperparams = { "n_units": "Number of neuronal cells to model in this layer", "batch_size": "Batch size dimension of this component", - "tau_m": "Cell membrane time constant", - "resist_m": "Membrane resistance value", + "tau_v": "Cell membrane time constant", "tau_w": "Recovery variable time constant", "v_reset": "Reset membrane potential value", "w_reset": "Reset angular driver value", @@ -229,11 +230,12 @@ def help(cls): ## component help function "omega": "Angular frequency of neuronal progress per second (radians)", "v0": "Initial condition for membrane potential/voltage", "w0": "Initial condition for membrane angular driver variable", + "resist_v": "Membrane resistance value", "integration_type": "Type of numerical integration to use for the cell dynamics" } info = {cls.__name__: properties, "compartments": compartment_props, - "dynamics": "tau_m * dv/dt = omega * w + v * b; " + "dynamics": "tau_v * dv/dt = omega * w + v * b; " "tau_w * dw/dt = w * b - v * omega + j", "hyperparameters": hyperparams} return info From 94f37f7e59d2c75e12afc3aa090d1a3c8ebec20a Mon Sep 17 00:00:00 2001 From: ago109 Date: Fri, 9 Aug 2024 12:38:05 -0400 Subject: [PATCH 26/61] cleaned up raf-cell --- ngclearn/components/neurons/spiking/RAFCell.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/ngclearn/components/neurons/spiking/RAFCell.py b/ngclearn/components/neurons/spiking/RAFCell.py index af49f00e..b550aee7 100755 --- a/ngclearn/components/neurons/spiking/RAFCell.py +++ b/ngclearn/components/neurons/spiking/RAFCell.py @@ -97,13 +97,13 @@ class RAFCell(JaxComponent): b: oscillation dampening factor (Default: -1) - v_reset: membrane potential reset condition (Default: 1 mV) + v_reset: reset condition for membrane potential (Default: 1 mV) w_reset: reset condition for angular current driver (Default: 0) - v0: membrane potential initial condition (Default: 1 mV) + v0: initial condition for membrane potential (Default: 1 mV) - w0: angular driver initial condition (Default: 0) + w0: initial condition for angular current driver (Default: 0) resist_v: membrane resistance (Default: 1 mega-Ohm) @@ -185,11 +185,11 @@ def advance_state(self, j, v, w, s, tols): self.tols.set(tols) @staticmethod - def _reset(batch_size, n_units, v_reset, w_reset): + def _reset(batch_size, n_units, v0, w0): restVals = jnp.zeros((batch_size, n_units)) j = restVals # None - v = restVals + v_reset - w = restVals + w_reset + v = restVals + v0 + w = restVals + w0 s = restVals #+ 0 tols = restVals #+ 0 return j, v, w, s, tols From 73e5aa14516e03119b2019eae918aab62a4683ae Mon Sep 17 00:00:00 2001 From: ago109 Date: Fri, 9 Aug 2024 12:46:54 -0400 Subject: [PATCH 27/61] cleaned up raf-cell --- ngclearn/components/neurons/spiking/RAFCell.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/ngclearn/components/neurons/spiking/RAFCell.py b/ngclearn/components/neurons/spiking/RAFCell.py index b550aee7..41949daf 100755 --- a/ngclearn/components/neurons/spiking/RAFCell.py +++ b/ngclearn/components/neurons/spiking/RAFCell.py @@ -159,14 +159,16 @@ def _advance_state(t, dt, tau_v, resist_v, tau_w, thr, omega, b, ## continue with centered dynamics j_ = j * resist_v if intgFlag == 1: ## RK-2/midpoint + ## Note: we integrate ODEs in order: first w, then v w_params = (j_, v, tau_w, omega, b) _, _w = step_rk2(0., w, _dfw, dt, w_params) - v_params = (j_, w, tau_v, omega, b) + v_params = (j_, _w, tau_v, omega, b) _, _v = step_rk2(0., v, _dfv, dt, v_params) else: # integType == 0 (default -- Euler) + ## Note: we integrate ODEs in order: first w, then v w_params = (j_, v, tau_w, omega, b) _, _w = step_euler(0., w, _dfw, dt, w_params) - v_params = (j_, w, tau_v, omega, b) + v_params = (j_, _w, tau_v, omega, b) _, _v = step_euler(0., v, _dfv, dt, v_params) s = _emit_spike(_v, thr) ## hyperpolarize/reset/snap variables From 6408ee0d51f2721ed53a9df937fb724b75d15f6d Mon Sep 17 00:00:00 2001 From: Alexander Ororbia Date: Sun, 11 Aug 2024 01:18:55 -0400 Subject: [PATCH 28/61] minor tweak to dim-reduce in utils --- ngclearn/utils/viz/dim_reduce.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/ngclearn/utils/viz/dim_reduce.py b/ngclearn/utils/viz/dim_reduce.py index 98646084..4f9095dd 100755 --- a/ngclearn/utils/viz/dim_reduce.py +++ b/ngclearn/utils/viz/dim_reduce.py @@ -1,6 +1,6 @@ import matplotlib import matplotlib.pyplot as plt -cmap = plt.cm.jet +default_cmap = plt.cm.jet import numpy as np from sklearn.decomposition import IncrementalPCA @@ -66,7 +66,8 @@ def extract_tsne_latents(vectors, perplexity=30, n_pca_comp=32): ## tSNE mapping z_2D = vectors return z_2D -def plot_latents(code_vectors, labels, plot_fname="2Dcode_plot.jpg", alpha=1.): +def plot_latents(code_vectors, labels, plot_fname="2Dcode_plot.jpg", alpha=1., + cmap=None): """ Produces a label-overlaid (label map to distinct colors) scatterplot for visualizing two-dimensional latent codes (produced by either PCA or t-SNE). @@ -80,7 +81,9 @@ def plot_latents(code_vectors, labels, plot_fname="2Dcode_plot.jpg", alpha=1.): plot_fname: /path/to/plot_fname. for saving the plot to disk - alpha: + alpha: alpha intensity level to present colors in scatterplot + + cmap: custom color-map to provide """ curr_backend = plt.rcParams["backend"] matplotlib.use('Agg') ## temporarily go in Agg plt backend for tsne plotting @@ -92,7 +95,11 @@ def plot_latents(code_vectors, labels, plot_fname="2Dcode_plot.jpg", alpha=1.): if lab.shape[1] > 1: ## extract integer class labels from a one-hot matrix lab = np.argmax(lab, 1) plt.figure(figsize=(8, 6)) - plt.scatter(code_vectors[:, 0], code_vectors[:, 1], c=lab, cmap=cmap, alpha=alpha) + _cmap = cmap + if _cmap is None: + _cmap = default_cmap + #print("> USING DEFAULT CMAP!") + plt.scatter(code_vectors[:, 0], code_vectors[:, 1], c=lab, cmap=_cmap, alpha=alpha) colorbar = plt.colorbar() #colorbar.set_alpha(1) #plt.draw_all() From 35eae76ea125ae7d38c7b1f293e340ca4f5a9151 Mon Sep 17 00:00:00 2001 From: Will Gebhardt Date: Tue, 19 Nov 2024 16:25:41 -0500 Subject: [PATCH 29/61] Additions for inhibition stuff --- ngclearn/components/base_monitor.py | 5 ++-- .../components/input_encoders/poissonCell.py | 2 +- ngclearn/components/other/varTrace.py | 24 +++++++++++-------- ngclearn/components/synapses/denseSynapse.py | 4 ++-- ngclearn/utils/viz/synapse_plot.py | 4 ++-- 5 files changed, 22 insertions(+), 17 deletions(-) diff --git a/ngclearn/components/base_monitor.py b/ngclearn/components/base_monitor.py index 8238264d..f3fae84e 100644 --- a/ngclearn/components/base_monitor.py +++ b/ngclearn/components/base_monitor.py @@ -124,9 +124,10 @@ def watch(self, compartment, window_length): """ cs, end = self._add_path(compartment.path) + dtype = compartment.value.dtype shape = compartment.value.shape - new_comp = Compartment(np.zeros(shape)) - new_comp_store = Compartment(np.zeros((window_length, *shape))) + new_comp = Compartment(np.zeros(shape, dtype=dtype)) + new_comp_store = Compartment(np.zeros((window_length, *shape), dtype=dtype)) comp_key = "*".join(compartment.path.split("/")) store_comp_key = comp_key + "*store" diff --git a/ngclearn/components/input_encoders/poissonCell.py b/ngclearn/components/input_encoders/poissonCell.py index 3c97fdb1..9128b8cd 100644 --- a/ngclearn/components/input_encoders/poissonCell.py +++ b/ngclearn/components/input_encoders/poissonCell.py @@ -46,7 +46,7 @@ class PoissonCell(JaxComponent): """ @deprecate_args(max_freq="target_freq") - def __init__(self, name, n_units, target_freq=0., batch_size=1, **kwargs): + def __init__(self, name, n_units, target_freq=63.75, batch_size=1, **kwargs): super().__init__(name, **kwargs) ## Constrained Bernoulli meta-parameters diff --git a/ngclearn/components/other/varTrace.py b/ngclearn/components/other/varTrace.py index 39dc2bb6..3c3eb625 100644 --- a/ngclearn/components/other/varTrace.py +++ b/ngclearn/components/other/varTrace.py @@ -5,7 +5,7 @@ from ngclearn.utils import tensorstats @partial(jit, static_argnums=[4]) -def _run_varfilter(dt, x, x_tr, decayFactor, a_delta=0.): +def _run_varfilter(dt, x, x_tr, decayFactor, gamma_tr, a_delta=0.): """ Run variable trace filter (low-pass filter) dynamics one step forward. @@ -22,7 +22,7 @@ def _run_varfilter(dt, x, x_tr, decayFactor, a_delta=0.): Returns: updated trace/filter value/state """ - _x_tr = x_tr * decayFactor + _x_tr = gamma_tr * x_tr * decayFactor #x_tr + (-x_tr) * (dt / tau_tr) = (1 - dt/tau_tr) * x_tr if a_delta > 0.: ## perform additive form of trace ODE _x_tr = _x_tr + x * a_delta @@ -64,13 +64,14 @@ class VarTrace(JaxComponent): ## low-pass filter """ # Define Functions - def __init__(self, name, n_units, tau_tr, a_delta, decay_type="exp", + def __init__(self, name, n_units, tau_tr, a_delta, 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.gamma_tr = gamma_tr self.decay_type = decay_type ## lin --> linear decay; exp --> exponential decay ## Layer Size Setup @@ -83,17 +84,20 @@ def __init__(self, name, n_units, tau_tr, a_delta, decay_type="exp", self.trace = Compartment(restVals) @staticmethod - def _advance_state(dt, decay_type, tau_tr, a_delta, inputs, trace): - ## compute the decay factor - decayFactor = 0. ## <-- pulse filter decay (default) + def _advance_state(dt, decay_type, tau_tr, a_delta, gamma_tr, inputs, trace): + decayFactor = 0. if "exp" in decay_type: decayFactor = jnp.exp(-dt/tau_tr) elif "lin" in decay_type: decayFactor = (1. - dt/tau_tr) - ## else "step" == decay_type, yielding a step/pulse-like filter - trace = _run_varfilter(dt, inputs, trace, decayFactor, a_delta) - outputs = trace - return outputs, trace + + _x_tr = gamma_tr * trace * decayFactor + if a_delta > 0.: + _x_tr = _x_tr + inputs * a_delta + else: + _x_tr = _x_tr * (1. - inputs) + inputs + + return trace, trace @resolver(_advance_state) def advance_state(self, outputs, trace): diff --git a/ngclearn/components/synapses/denseSynapse.py b/ngclearn/components/synapses/denseSynapse.py index 957689c6..2988ac90 100755 --- a/ngclearn/components/synapses/denseSynapse.py +++ b/ngclearn/components/synapses/denseSynapse.py @@ -47,8 +47,8 @@ def __init__(self, name, shape, weight_init=None, bias_init=None, self.bias_init = bias_init ## Synapse meta-parameters - self.shape = shape ## shape of synaptic efficacy matrix - self.Rscale = resist_scale ## post-transformation scale factor + self.shape = shape + self.Rscale = resist_scale ## Set up synaptic weight values tmp_key, *subkeys = random.split(self.key.value, 4) diff --git a/ngclearn/utils/viz/synapse_plot.py b/ngclearn/utils/viz/synapse_plot.py index 9c4a0e49..14912907 100644 --- a/ngclearn/utils/viz/synapse_plot.py +++ b/ngclearn/utils/viz/synapse_plot.py @@ -138,13 +138,13 @@ def visualize_gif(frames, path='.', name='tmp', suffix='.jpg', **kwargs): _frames = [f.astype(jnp.uint8) for f in frames] iio.imwrite(path + '/' + name + '.gif', _frames, **kwargs) -def make_video(f_start, f_end, path, prefix, suffix='.jpg', skip=1): +def make_video(f_start, f_end, path, prefix, suffix='.jpg', skip=1, **kwargs): images = [] for i in range(f_start, f_end+1, skip): print("Reading frame " + str(i)) images.append(iio.imread(path + "/" + prefix + str(i) + suffix)) print("writing gif") - iio.imwrite(path + '/training.gif', images, loop=0, duration=200) + iio.imwrite(path + '/training.gif', images, **kwargs) # def visualize_norm(thetas, sizes, prefix, suffix='.jpg'): From 796178d582f42989f57bf6b0d0f31c3dff80e0ea Mon Sep 17 00:00:00 2001 From: Alexander Ororbia Date: Fri, 28 Feb 2025 20:01:49 -0500 Subject: [PATCH 30/61] commit probes/mods to utils to analysis_tools branch --- ngclearn/utils/analysis/__init__.py | 3 + ngclearn/utils/analysis/attentive_probe.py | 170 +++++++++++++++++++++ ngclearn/utils/analysis/linear_probe.py | 125 +++++++++++++++ ngclearn/utils/analysis/probe.py | 84 ++++++++++ ngclearn/utils/model_utils.py | 42 ++++- 5 files changed, 423 insertions(+), 1 deletion(-) create mode 100644 ngclearn/utils/analysis/__init__.py create mode 100644 ngclearn/utils/analysis/attentive_probe.py create mode 100644 ngclearn/utils/analysis/linear_probe.py create mode 100644 ngclearn/utils/analysis/probe.py diff --git a/ngclearn/utils/analysis/__init__.py b/ngclearn/utils/analysis/__init__.py new file mode 100644 index 00000000..0f97d6e4 --- /dev/null +++ b/ngclearn/utils/analysis/__init__.py @@ -0,0 +1,3 @@ +## point to supported analysis probes +from .linear_probe import LinearProbe +from .attentive_probe import AttentiveProbe diff --git a/ngclearn/utils/analysis/attentive_probe.py b/ngclearn/utils/analysis/attentive_probe.py new file mode 100644 index 00000000..115c0e1c --- /dev/null +++ b/ngclearn/utils/analysis/attentive_probe.py @@ -0,0 +1,170 @@ +import jax +import numpy as np +from ngclearn.utils.analysis.probe import Probe +from ngclearn.utils.model_utils import drop_out, softmax, gelu, layer_normalize +from ngclearn.utils.optim import adam +from jax import jit, random, numpy as jnp, lax, nn +from functools import partial as bind + +def masked_fill(x: jax.Array, mask: jax.Array, value=0) -> jax.Array: + """ + Return an output with masked condition, with non-masked value + be the other value + + Args: + x (jax.Array): _description_ + mask (jax.Array): _description_ + value (int, optional): _description_. Defaults to 0. + + Returns: + jax.Array: _description_ + """ + return jnp.where(mask, jnp.broadcast_to(value, x.shape), x) + +@bind(jax.jit, static_argnums=[4, 5]) +def cross_attention(params: tuple, x1: jax.Array, x2: jax.Array, mask: jax.Array, n_heads: int=8, dropout_rate: float=0.0): + B, T, Dq = x1.shape # The original shape + _, S, Dkv = x2.shape + # in here we attend x2 to x1 + Wq, bq, Wk, bk, Wv, bv, Wout, bout = params + # projection + q = x1 @ Wq + bq # normal linear transformation (B, T, D) + k = x2 @ Wk + bk # normal linear transformation (B, S, D) + v = x2 @ Wv + bv # normal linear transformation (B, S, D) + hidden = q.shape[-1] + _hidden = hidden // n_heads + q = q.reshape((B, T, n_heads, _hidden)).transpose([0, 2, 1, 3]) # (B, H, T, D) + k = k.reshape((B, S, n_heads, _hidden)).transpose([0, 2, 1, 3]) # (B, H, T, D) + v = v.reshape((B, S, n_heads, _hidden)).transpose([0, 2, 1, 3]) # (B, H, T, D) + score = jnp.einsum("BHTE,BHSE->BHTS", q, k) / jnp.sqrt(_hidden) # Q @ KT / ||d||; d = D // n_heads + if mask is not None: + Tq, Tk = q.shape[2], k.shape[2] + assert mask.shape == (B, Tq, Tk), (mask.shape, (B, Tq, Tk)) + _mask = mask.reshape((B, 1, Tq, Tk)) # 'b tq tk -> b 1 tq tk' + score = masked_fill(score, _mask, value=-jnp.inf) # basically masking out all must-unattended values + score = jax.nn.softmax(score, axis=-1) # (B, H, T, S) + score = score.astype(q.dtype) # (B, H, T, S) + if dropout_rate > 0.: + score = drop_out(input=score, rate=dropout_rate) ## NOTE: normally you apply dropout here + attention = jnp.einsum("BHTS,BHSE->BHTE", score, v) # (B, T, H, E) + attention = attention.transpose([0, 2, 1, 3]).reshape((B, T, -1)) # (B, T, H, E) => (B, T, D) + return attention @ Wout + bout # (B, T, Dq) + +@bind(jax.jit, static_argnums=[3, 4, 5, 6]) +def run_attention_probe(params, encodings, mask, n_heads: int, dropout: float = 0.0, use_LN=False, use_softmax=True): + # encoded_image_feature: (B, hw, dim) + #learnable_query, *_params) = params + learnable_query, Wq, bq, Wk, bk, Wv, bv, Wout, bout, Whid, bhid, Wln_mu, Wln_scale, Wy, by = params + attn_params = (Wq, bq, Wk, bk, Wv, bv, Wout, bout) + features = cross_attention(attn_params, learnable_query, encodings, mask, n_heads, dropout) + features = features[:, 0] # (B, 1, dim) => (B, dim) + hids = jnp.matmul((features + learnable_query[:, 0]), Whid) + bhid + hids = gelu(hids) + if use_LN: ## normalize hidden layer output of probe predictor + hids = layer_normalize(hids, Wln_mu, Wln_scale) + outs = jnp.matmul(hids, Wy) + by + if use_softmax: ## apply softmax output nonlinearity + outs = softmax(outs) + return outs, features + +@bind(jax.jit, static_argnums=[4, 5, 6, 7]) +def eval_attention_probe(params, encodings, labels, mask, n_heads: int, dropout: float = 0.0, use_LN=False, use_softmax=True): + # encodings: (B, hw, dim) + outs, _ = run_attention_probe(params, encodings, mask, n_heads, dropout, use_LN, use_softmax) + if use_softmax: ## Multinoulli log likelihood for 1-of-K predictions + L = -jnp.mean(jnp.sum(jnp.log(outs) * labels, axis=1, keepdims=True)) + else: ## MSE for real-valued outputs + L = jnp.mean(jnp.sum(jnp.square(outs - labels), axis=1, keepdims=True)) + return L, outs #, features + +class AttentiveProbe(Probe): + """ + Args: + dkey: init seed key + + source_seq_length: length of input sequence (e.g., height x width of the image feature) + + input_dim: input dimensionality of probe + + out_dim: output dimensionality of probe + + num_heads: number of cross-attention heads + + head_dim: output dimensionality of each cross-attention head + + target_seq_length: to pool, we set it at one (or map the source sequence to the target sequence of length 1) + + learnable_query_dim: target sequence dim (output dimension of cross-attention portion of probe) + + batch_size: size of batches to process per internal call to update (or process) + + hid_dim: dimensionality of hidden layer(s) of MLP portion of probe + + use_LN: should layer normalization be used within MLP portions of probe or not? + + use_softmax: should a softmax be applied to output of probe or not? + + """ + def __init__( + self, dkey, source_seq_length, input_dim, out_dim, num_heads=8, head_dim=64, + target_seq_length=1, learnable_query_dim=31, batch_size=1, hid_dim=32, use_LN=True, use_softmax=True, **kwargs + ): + super().__init__(dkey, batch_size, **kwargs) + self.dkey, *subkeys = random.split(self.dkey, 12) + self.num_heads = num_heads + self.source_seq_length = source_seq_length + self.input_dim = input_dim + self.out_dim = out_dim + self.use_softmax = use_softmax + self.use_LN = use_LN + + sigma = 0.05 + ## cross-attention parameters + Wq = random.normal(subkeys[0], (learnable_query_dim, head_dim)) * sigma + bq = random.normal(subkeys[1], (1, head_dim)) * sigma + Wk = random.normal(subkeys[2], (input_dim, head_dim)) * sigma + bk = random.normal(subkeys[3], (1, head_dim)) * sigma + Wv = random.normal(subkeys[4], (input_dim, head_dim)) * sigma + bv = random.normal(subkeys[5], (1, head_dim)) * sigma + Wout = random.normal(subkeys[6], (head_dim, learnable_query_dim)) * sigma + bout = random.normal(subkeys[7], (1, learnable_query_dim)) * sigma + #params = (Wq, bq, Wk, bk, Wv, bv, Wout, bout) + learnable_query = jnp.zeros((batch_size, 1, learnable_query_dim)) # (B, T, D) + #self.all_params = (learnable_query, *params) + self.mask = np.zeros((batch_size, target_seq_length, source_seq_length)).astype(bool) ## mask tensor + ## MLP parameters + Whid = random.normal(subkeys[8], (learnable_query_dim, hid_dim)) * sigma + bhid = random.normal(subkeys[9], (1, hid_dim)) * sigma + Wln_mu = jnp.zeros((1, hid_dim)) + Wln_scale = jnp.ones((1, hid_dim)) + Wy = random.normal(subkeys[8], (hid_dim, out_dim)) * sigma + by = random.normal(subkeys[9], (1, out_dim)) * sigma + #mlp_params = (Whid, bhid, Wln_mu, Wln_scale, Wy, by) + self.probe_params = (learnable_query, Wq, bq, Wk, bk, Wv, bv, Wout, bout, Whid, bhid, Wln_mu, Wln_scale, Wy, by) + + ## set up gradient calculator + self.grad_fx = jax.value_and_grad(eval_attention_probe, argnums=0, has_aux=True) + ## set up update rule/optimizer + self.optim_params = adam.adam_init(self.probe_params) + self.eta = 0.001 + + def process(self, embedding_sequence): + outs, feats = run_attention_probe( + self.probe_params, embedding_sequence, self.mask, self.num_heads, 0.0, use_LN=self.use_LN, + use_softmax=self.use_softmax + ) + return outs + + def update(self, embedding_sequence, labels): + ## compute partial derivatives / adjustments to probe parameters + outputs, grads = self.grad_fx( + self.probe_params, embedding_sequence, labels, self.mask, self.num_heads, dropout=0., use_LN=self.use_LN, + use_softmax=self.use_softmax + ) + loss, predictions = outputs + ## adjust parameters of probe + self.optim_params, self.probe_params = adam.adam_step( + self.optim_params, self.probe_params, grads, eta=self.eta + ) + return loss, predictions + diff --git a/ngclearn/utils/analysis/linear_probe.py b/ngclearn/utils/analysis/linear_probe.py new file mode 100644 index 00000000..86eeaf5c --- /dev/null +++ b/ngclearn/utils/analysis/linear_probe.py @@ -0,0 +1,125 @@ +import jax +import numpy as np +from ngclearn.utils.analysis.probe import Probe +from ngclearn.utils.model_utils import drop_out, softmax, layer_normalize +from jax import jit, random, numpy as jnp, lax, nn +from functools import partial as bind +import ngclearn.utils.weight_distribution as dist +from ngclearn.utils.optim import adam, sgd + +@bind(jax.jit, static_argnums=[2, 3]) +def run_linear_probe(params, x, use_softmax=False, use_LN=False): + Wln_mu, Wln_scale, W, b = params + _x = x + if use_LN: ## normalize input vector to probe predictor + _x = layer_normalize(_x, Wln_mu, Wln_scale) + y_mu = (jnp.matmul(_x, W) + b) + if use_softmax: + y_mu = softmax(y_mu) + return y_mu + +@bind(jax.jit, static_argnums=[3, 4]) +def eval_linear_probe(params, x, y, use_softmax=True, use_LN=False): + y_mu = run_linear_probe(params, x, use_softmax=use_softmax, use_LN=use_LN) + e = y_mu - y + if use_softmax: ## Multinoulli log likelihood for 1-of-K predictions + L = -jnp.mean(jnp.sum(jnp.log(y_mu) * y, axis=1, keepdims=True)) + else: ## MSE for real-valued outputs + L = jnp.sum(jnp.square(e)) * 1./x.shape[0] + return L, y_mu + #return y_mu, L, e + +# @bind(jax.jit, static_argnums=[6, 7]) +# def calc_linear_probe_grad(x, y, params, eta, decay=0., l1_decay=0., use_softmax=False, use_LN=False): +# y_mu, L, e = eval_linear_probe(params, x, y, use_softmax=use_softmax, use_LN=use_LN) +# Wln_mu, Wln_scale, W, b = params +# dW = jnp.matmul(x.T, e) + W * decay/eta + jnp.abs(W) * 0.5 * l1_decay/eta +# db = jnp.sum(e, axis=0, keepdims=True) +# dW = dW * (1. / x.shape[0]) +# db = db * (1. / x.shape[0]) +# return y_mu, L, [dW, db] + +# @jit +# def update_linear_probe(x, y, params, eta, decay=0., l1_decay=0., use_softmax=False): +# y_mu, L, e = run_linear_probe(x, params, use_softmax=use_softmax) +# W, b = params +# dW = jnp.matmul(x.T, e) +# db = jnp.sum(e, axis=0, keepdims=True) +# W = W - dW * eta/x.shape[0] - W * decay/x.shape[0] - jnp.abs(W) * 0.5 * l1_decay/x.shape[0] +# b = b - db * eta/x.shape[0] +# return y_mu, L, [W, b] + +class LinearProbe(Probe): + """ + Args: + dkey: init seed key + + source_seq_length: length of input sequence (e.g., height x width of the image feature) + + input_dim: input dimensionality of probe + + out_dim: output dimensionality of probe + + batch_size: size of batches to process per internal call to update (or process) + + use_LN: should layer normalization be used on incoming input vectors given to this probe? + + use_softmax: should a softmax be applied to output of probe or not? + + """ + def __init__( + self, dkey, source_seq_length, input_dim, out_dim, batch_size=1, use_LN=False, use_softmax=False, **kwargs + ): + super().__init__(dkey, batch_size, **kwargs) + self.dkey, *subkeys = random.split(self.dkey, 3) + self.source_seq_length = source_seq_length + self.input_dim = input_dim + self.out_dim = out_dim + self.use_softmax = use_softmax + self.use_LN = use_LN + self.l2_decay = 0.0001 + self.l1_decay = 0.000025 + ## TODO: add in pre-built layer norm of inputs? + + ## set up classifier + flat_input_dim = input_dim * source_seq_length + weight_init = dist.fan_in_gaussian() # dist.gaussian(mu=0., sigma=0.05) # 0.02) + Wln_mu = jnp.zeros((1, flat_input_dim)) + Wln_scale = jnp.ones((1, flat_input_dim)) + W = dist.initialize_params(subkeys[0], weight_init, (flat_input_dim, out_dim)) + b = jnp.zeros((1, out_dim)) + self.probe_params = [Wln_mu, Wln_scale, W, b] + + ## set up update rule/optimizer + ## set up gradient calculator + self.grad_fx = jax.value_and_grad(eval_linear_probe, argnums=0, has_aux=True) + self.optim_params = adam.adam_init(self.probe_params) + self.eta = 0.001 + + def process(self, embeddings): + _embeddings = embeddings + if len(_embeddings.shape) > 2: + flat_dim = embeddings.shape[1] * embeddings.shape[2] + _embeddings = jnp.reshape(_embeddings, (embeddings.shape[0], flat_dim)) + outs = run_linear_probe(self.probe_params, _embeddings, use_softmax=self.use_softmax, use_LN=self.use_LN) + return outs + + def update(self, embeddings, labels): + _embeddings = embeddings + if len(_embeddings.shape) > 2: + flat_dim = embeddings.shape[1] * embeddings.shape[2] + _embeddings = jnp.reshape(_embeddings, (embeddings.shape[0], flat_dim)) + ## compute adjustments to probe parameters + # predictions, loss, grads = calc_linear_probe_grad( + # self.probe_params, _embeddings, labels, self.eta, decay=self.l2_decay, l1_decay=self.l1_decay, + # use_softmax=self.use_softmax, use_LN=self.use_LN + # ) + outputs, grads = self.grad_fx( + self.probe_params, _embeddings, labels, use_softmax=self.use_softmax, use_LN=self.use_LN + ) + loss, predictions = outputs + ## adjust parameters of probe + self.optim_params, self.probe_params = adam.adam_step( + self.optim_params, self.probe_params, grads, eta=self.eta + ) + return loss, predictions diff --git a/ngclearn/utils/analysis/probe.py b/ngclearn/utils/analysis/probe.py new file mode 100644 index 00000000..652ae90c --- /dev/null +++ b/ngclearn/utils/analysis/probe.py @@ -0,0 +1,84 @@ +from jax import random, numpy as jnp + +class Probe(): + """ + General framework for an analysis probe (that may or may not be learnable in an iterative fashion). + + Args: + dkey: init seed key + + batch_size: size of batches to process per internal call to update (or process) + + """ + def __init__( + self, dkey, batch_size=4, **kwargs + ): + #dkey, *subkeys = random.split(dkey, 3) + self.dkey = dkey + self.batch_size = batch_size + + def process(self, embeddings): + predictions = None + return predictions + + def update(self, embeddings, labels): + L = predictions = None + return L, predictions + + def predict(self, data): + _data = data + if len(_data.shape) < 3: + _data = jnp.expand_dims(_data, axis=1) + + n_samples, seq_len, dim = _data.shape + n_batches = int(n_samples / self.batch_size) + s_ptr = 0 + e_ptr = self.batch_size + Y_mu = [] + for b in range(n_batches): + x_mb = _data[s_ptr:e_ptr, :, :] ## slice out 3D batch tensor + s_ptr = e_ptr + e_ptr += x_mb.shape[0] + y_mu = self.process(x_mb) + Y_mu.append(y_mu) + Y_mu = jnp.concatenate(Y_mu, axis=0) + return Y_mu + + def fit(self, data, labels, n_iter=50): + _data = data + if len(_data.shape) < 3: + _data = jnp.expand_dims(_data, axis=1) + + n_samples, seq_len, dim = _data.shape + n_batches = int(n_samples / self.batch_size) + + Y_mu = [] + _Y = None + for iter in range(n_iter): + ## shuffle data (to ensure i.i.d. across sequences) + self.dkey, *subkeys = random.split(self.dkey, 2) + ptrs = random.permutation(subkeys[0], n_samples) + _X = _data[ptrs, :, :] + _Y = labels[ptrs, :] + ## run one epoch over data tensors + L = 0. + Ns = 0. + + s_ptr = 0 + e_ptr = self.batch_size + for b in range(n_batches): + x_mb = _X[s_ptr:e_ptr, :, :] ## slice out 3D batch tensor + y_mb = _Y[s_ptr:e_ptr, :] + s_ptr = e_ptr + e_ptr += x_mb.shape[0] + Ns += x_mb.shape[0] + + _L, py = self.update(x_mb, y_mb) + L = _L + L + print(f"\r{iter} L = {L/Ns}", end="") # p(y|z):\n{py}") + if iter == n_iter-1: + Y_mu.append(py) + print() + if iter == n_iter - 1: + Y_mu = jnp.concatenate(Y_mu, axis=0) + return Y_mu, _Y diff --git a/ngclearn/utils/model_utils.py b/ngclearn/utils/model_utils.py index 59bc8a32..140ad542 100755 --- a/ngclearn/utils/model_utils.py +++ b/ngclearn/utils/model_utils.py @@ -83,6 +83,9 @@ def create_function(fun_name, args=None): if fun_name == "tanh": fx = tanh dfx = d_tanh + elif "kwta" in fun_name: + fx = bkwta + dfx = bkwta #d_identity elif fun_name == "sigmoid": fx = sigmoid dfx = d_sigmoid @@ -98,6 +101,9 @@ def create_function(fun_name, args=None): elif fun_name == "softplus": fx = softplus dfx = d_softplus + elif fun_name == "softmax": + fx = softmax + dfx = d_identity ## TODO: currently Jacobian of softmax not supported! elif fun_name == "unit_threshold": fx = threshold ## default threshold is 1 (thus unit) dfx = d_threshold ## STE approximation @@ -109,10 +115,17 @@ def create_function(fun_name, args=None): dfx = d_identity else: raise RuntimeError( - "Activition function (" + fun_name + ") is not recognized/supported!" + "Activation function (" + fun_name + ") is not recognized/supported!" ) return fx, dfx +@partial(jit, static_argnums=[1]) +def bkwta(x, nWTA=5): #5 10 15 #K=50): + values, indices = lax.top_k(x, nWTA) # Note: we do not care to sort the indices + kth = jnp.expand_dims(jnp.min(values,axis=1),axis=1) # must do comparison per sample in potential mini-batch + topK = jnp.greater_equal(x, kth).astype(jnp.float32) # cast booleans to floats + return topK + @partial(jit, static_argnums=[2, 3, 4]) def normalize_matrix(M, wnorm, order=1, axis=0, scale=1.): """ @@ -440,6 +453,26 @@ def inverse_logistic(x, clip_bound=0.03): # 0.03 x_ = jnp.clip(x_, clip_bound, 1.0 - clip_bound) return jnp.log( x_/((1.0 - x_) + 1e-6) ) +@jit +def gelu(x): + """ + Applies the Gaussian Error Linear Unit (GeLU) activation (specifically, a fast approximation is used). + + Args: + x: data to transform via inverse logistic function + + Returns: + output of the GeLU activation + """ + return x * sigmoid(x * 1.702) ## approximate GeLU + +@jit +def d_gelu(x): + # df/dx = 1.702 * [ 1/(exp(-x) + 1) + (exp(-x) * x) / (exp(-x) + 1)^2] + exp_neg_x = jnp.exp(-x) + _x = (1./(exp_neg_x + 1.)) + (exp_neg_x * x)/jnp.square(exp_neg_x+1) + return _x * 1.702 + @jit def softmax(x, tau=0.0): """ @@ -499,6 +532,13 @@ def threshold_cauchy(x, lmbda): term2 = g * (x <= -lmbda).astype(jnp.float32) ## g * (x <= -lmda) return term1 + term2 +@jit +def layer_normalize(x, shift=0., scale=1.): + xmu = jnp.mean(x, axis=1, keepdims=True) + xsigma = jnp.sqrt(jnp.mean(jnp.square(x - xmu))) + _x = (x - xmu)/(xsigma + 1e-6) + return _x * scale + shift + @jit def drop_out(dkey, input, rate=0.0): """ From 84237ffb2d9382c3a228ed9f7c53b11022db5353 Mon Sep 17 00:00:00 2001 From: Alexander Ororbia Date: Fri, 28 Feb 2025 20:05:51 -0500 Subject: [PATCH 31/61] commit probes/mods to utils to analysis_tools branch --- ngclearn/utils/model_utils.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/ngclearn/utils/model_utils.py b/ngclearn/utils/model_utils.py index 140ad542..02adc01b 100755 --- a/ngclearn/utils/model_utils.py +++ b/ngclearn/utils/model_utils.py @@ -534,6 +534,19 @@ def threshold_cauchy(x, lmbda): @jit def layer_normalize(x, shift=0., scale=1.): + """ + Applies layer normalization to input data `x` + + Args: + x: data to apply threshold function over + + shift: the compensating mean/shift factor/parameters (to undo mean subtraction) + + scale: the compensating re-scaling factor/parameters (to undo standard deviation division) + + Returns: + layer-normalized data samples `x` + """ xmu = jnp.mean(x, axis=1, keepdims=True) xsigma = jnp.sqrt(jnp.mean(jnp.square(x - xmu))) _x = (x - xmu)/(xsigma + 1e-6) From 9d7acbb28878702304477e4b4fbeb453621fecc8 Mon Sep 17 00:00:00 2001 From: Viet Dung Nguyen Date: Sat, 1 Mar 2025 13:05:04 -0500 Subject: [PATCH 32/61] update documentation --- ngclearn/utils/analysis/attentive_probe.py | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/ngclearn/utils/analysis/attentive_probe.py b/ngclearn/utils/analysis/attentive_probe.py index 115c0e1c..7790868d 100644 --- a/ngclearn/utils/analysis/attentive_probe.py +++ b/ngclearn/utils/analysis/attentive_probe.py @@ -22,7 +22,25 @@ def masked_fill(x: jax.Array, mask: jax.Array, value=0) -> jax.Array: return jnp.where(mask, jnp.broadcast_to(value, x.shape), x) @bind(jax.jit, static_argnums=[4, 5]) -def cross_attention(params: tuple, x1: jax.Array, x2: jax.Array, mask: jax.Array, n_heads: int=8, dropout_rate: float=0.0): +def cross_attention(params: tuple, x1: jax.Array, x2: jax.Array, mask: jax.Array, n_heads: int=8, dropout_rate: float=0.0) -> jax.Array: + """ + Run cross-attention function given a list of parameters and two sequences (x1 and x2). + The function takes in a query sequence x1 and a key-value sequence x2, and returns an output of the same shape as x1. + T is the length of the query sequence, and S is the length of the key-value sequence. + Dq is the dimension of the query sequence, and Dkv is the dimension of the key-value sequence. + H is the number of attention heads. + + Args: + params (tuple): tuple of parameters + x1 (jax.Array): query sequence. Shape: (B, T, Dq) + x2 (jax.Array): key-value sequence. Shape: (B, S, Dkv) + mask (jax.Array): mask tensor. Shape: (B, T, S) + n_heads (int, optional): number of attention heads. Defaults to 8. + dropout_rate (float, optional): dropout rate. Defaults to 0.0. + + Returns: + jax.Array: output of cross-attention + """ B, T, Dq = x1.shape # The original shape _, S, Dkv = x2.shape # in here we attend x2 to x1 From 247de74e357da22ed88e1e6b6546368232b939b6 Mon Sep 17 00:00:00 2001 From: Alexander Ororbia Date: Sat, 1 Mar 2025 14:24:09 -0500 Subject: [PATCH 33/61] cleaned up probes/docs for probes --- ngclearn/utils/analysis/attentive_probe.py | 75 ++++++++++++++++++++++ ngclearn/utils/analysis/linear_probe.py | 30 ++++++++- ngclearn/utils/analysis/probe.py | 30 ++++++++- 3 files changed, 132 insertions(+), 3 deletions(-) diff --git a/ngclearn/utils/analysis/attentive_probe.py b/ngclearn/utils/analysis/attentive_probe.py index 7790868d..e919eb49 100644 --- a/ngclearn/utils/analysis/attentive_probe.py +++ b/ngclearn/utils/analysis/attentive_probe.py @@ -32,10 +32,15 @@ def cross_attention(params: tuple, x1: jax.Array, x2: jax.Array, mask: jax.Array Args: params (tuple): tuple of parameters + x1 (jax.Array): query sequence. Shape: (B, T, Dq) + x2 (jax.Array): key-value sequence. Shape: (B, S, Dkv) + mask (jax.Array): mask tensor. Shape: (B, T, S) + n_heads (int, optional): number of attention heads. Defaults to 8. + dropout_rate (float, optional): dropout rate. Defaults to 0.0. Returns: @@ -70,6 +75,27 @@ def cross_attention(params: tuple, x1: jax.Array, x2: jax.Array, mask: jax.Array @bind(jax.jit, static_argnums=[3, 4, 5, 6]) def run_attention_probe(params, encodings, mask, n_heads: int, dropout: float = 0.0, use_LN=False, use_softmax=True): + """ + Runs full nonlinear attentive probe on input encodings (typically embedding vectors produced by some other model). + + Args: + params: parameters tuple/list of probe + + encodings: input encoding vectors/data + + mask: optional mask to be applied to internal cross-attention + + n_heads: number of attention heads + + dropout: if >0, triggers drop-out applied internally to cross-attention + + use_LN: use layer normalization? + + use_softmax: should softmax be applied to output of attention probe? (useful for classification) + + Returns: + output scores/probabilities, cross-attention (hidden) features + """ # encoded_image_feature: (B, hw, dim) #learnable_query, *_params) = params learnable_query, Wq, bq, Wk, bk, Wv, bv, Wout, bout, Whid, bhid, Wln_mu, Wln_scale, Wy, by = params @@ -87,6 +113,30 @@ def run_attention_probe(params, encodings, mask, n_heads: int, dropout: float = @bind(jax.jit, static_argnums=[4, 5, 6, 7]) def eval_attention_probe(params, encodings, labels, mask, n_heads: int, dropout: float = 0.0, use_LN=False, use_softmax=True): + """ + Runs and evaluates the nonlinear attentive probe given a paired set of encoding vectors and externally assigned + labels/regression targets. + + Args: + params: parameters tuple/list of probe + + encodings: input encoding vectors/data + + labels: output target values (e.g., labels, regression target vectors) + + mask: optional mask to be applied to internal cross-attention + + n_heads: number of attention heads + + dropout: if >0, triggers drop-out applied internally to cross-attention + + use_LN: use layer normalization? + + use_softmax: should softmax be applied to output of attention probe? (useful for classification) + + Returns: + current loss value, output scores/probabilities + """ # encodings: (B, hw, dim) outs, _ = run_attention_probe(params, encodings, mask, n_heads, dropout, use_LN, use_softmax) if use_softmax: ## Multinoulli log likelihood for 1-of-K predictions @@ -97,6 +147,10 @@ def eval_attention_probe(params, encodings, labels, mask, n_heads: int, dropout: class AttentiveProbe(Probe): """ + This implements a nonlinear attentive probe, which is useful for evaluating the quality of + encodings/embeddings in light of some superivsory downstream data (e.g., label one-hot + encodings or real-valued vector regression targets). + Args: dkey: init seed key @@ -167,6 +221,15 @@ def __init__( self.eta = 0.001 def process(self, embedding_sequence): + """ + Runs the probe's inference scheme given an input batch of sequences of encodings/embeddings. + + Args: + embedding_sequence: a 3D tensor containing a batch of encoding sequences; shape (B, T, embed_dim) + + Returns: + probe output scores/probability values + """ outs, feats = run_attention_probe( self.probe_params, embedding_sequence, self.mask, self.num_heads, 0.0, use_LN=self.use_LN, use_softmax=self.use_softmax @@ -174,6 +237,18 @@ def process(self, embedding_sequence): return outs def update(self, embedding_sequence, labels): + """ + Runs and updates this probe given an input batch of sequences of encodings/embeddings and their externally + assigned labels/target vector values. + + Args: + embedding_sequence: a 3D tensor containing a batch of encoding sequences; shape (B, T, embed_dim) + + labels: target values that map to embedding sequence; shape (B, target_value_dim) + + Returns: + probe output scores/probability values + """ ## compute partial derivatives / adjustments to probe parameters outputs, grads = self.grad_fx( self.probe_params, embedding_sequence, labels, self.mask, self.num_heads, dropout=0., use_LN=self.use_LN, diff --git a/ngclearn/utils/analysis/linear_probe.py b/ngclearn/utils/analysis/linear_probe.py index 86eeaf5c..05284e94 100644 --- a/ngclearn/utils/analysis/linear_probe.py +++ b/ngclearn/utils/analysis/linear_probe.py @@ -51,6 +51,11 @@ def eval_linear_probe(params, x, y, use_softmax=True, use_LN=False): class LinearProbe(Probe): """ + This implements a regularized linear probe, which is useful for evaluating the quality of + encodings/embeddings in light of some superivsory downstream data (e.g., label one-hot + encodings or real-valued vector regression targets). + Note that this probe allows for configurable Elastic-net (L1+L2) regularization. + Args: dkey: init seed key @@ -79,7 +84,6 @@ def __init__( self.use_LN = use_LN self.l2_decay = 0.0001 self.l1_decay = 0.000025 - ## TODO: add in pre-built layer norm of inputs? ## set up classifier flat_input_dim = input_dim * source_seq_length @@ -97,14 +101,35 @@ def __init__( self.eta = 0.001 def process(self, embeddings): + """ + Runs the probe's inference scheme given an input batch of sequences of encodings/embeddings. + + Args: + embedding_sequence: a 3D tensor containing a batch of encoding sequences; shape (B, T, embed_dim) + + Returns: + probe output scores/probability values + """ _embeddings = embeddings - if len(_embeddings.shape) > 2: + if len(_embeddings.shape) > 2: ## we flatten a sequence batch to 2D for a linear probe flat_dim = embeddings.shape[1] * embeddings.shape[2] _embeddings = jnp.reshape(_embeddings, (embeddings.shape[0], flat_dim)) outs = run_linear_probe(self.probe_params, _embeddings, use_softmax=self.use_softmax, use_LN=self.use_LN) return outs def update(self, embeddings, labels): + """ + Runs and updates this probe given an input batch of sequences of encodings/embeddings and their externally + assigned labels/target vector values. + + Args: + embedding_sequence: a 3D tensor containing a batch of encoding sequences; shape (B, T, embed_dim) + + labels: target values that map to embedding sequence; shape (B, target_value_dim) + + Returns: + probe output scores/probability values + """ _embeddings = embeddings if len(_embeddings.shape) > 2: flat_dim = embeddings.shape[1] * embeddings.shape[2] @@ -123,3 +148,4 @@ def update(self, embeddings, labels): self.optim_params, self.probe_params, grads, eta=self.eta ) return loss, predictions + diff --git a/ngclearn/utils/analysis/probe.py b/ngclearn/utils/analysis/probe.py index 652ae90c..d9aa1cf2 100644 --- a/ngclearn/utils/analysis/probe.py +++ b/ngclearn/utils/analysis/probe.py @@ -26,6 +26,15 @@ def update(self, embeddings, labels): return L, predictions def predict(self, data): + """ + Runs this probe's inference scheme over a pool of data. + + Args: + data: a dataset or design tensor/matrix containing encoding vector sequences; shape (N, T, embed_dim) or (N, embed_dim) + + Returns: + the output scores/predictions made by this probe + """ _data = data if len(_data.shape) < 3: _data = jnp.expand_dims(_data, axis=1) @@ -45,13 +54,31 @@ def predict(self, data): return Y_mu def fit(self, data, labels, n_iter=50): + """ + Fits this probe to a pool of data. + + Args: + data: a dataset or design tensor/matrix containing encoding vector sequences; shape (N, T, embed_dim) or (N, embed_dim) + + labels: a design matrix containing corresponding labels/targets for the embedding data; shape (N, target_dim) + + Returns: + the output scores/predictions made by this probe + """ _data = data if len(_data.shape) < 3: _data = jnp.expand_dims(_data, axis=1) n_samples, seq_len, dim = _data.shape + size_modulo = n_samples % self.batch_size + if size_modulo > 0: + ## we append some dup data for dataset design tensors that do not divide by batch size evenly + _chunk = _data[0:size_modulo, :, :] + _data = jnp.concatenate((_data, _chunk), axis=0) + n_samples, seq_len, dim = _data.shape n_batches = int(n_samples / self.batch_size) + ## run main probe fitting loop Y_mu = [] _Y = None for iter in range(n_iter): @@ -81,4 +108,5 @@ def fit(self, data, labels, n_iter=50): print() if iter == n_iter - 1: Y_mu = jnp.concatenate(Y_mu, axis=0) - return Y_mu, _Y + return Y_mu, _Y ## return predictions mapped to current shuffling of labels + From d0df86ee5a883e7a3afc028cdd25edcac3a38503 Mon Sep 17 00:00:00 2001 From: Viet Dung Nguyen Date: Sat, 1 Mar 2025 15:14:45 -0500 Subject: [PATCH 34/61] change heads_dim to attn_dim, and modify the mlp to be as similar as possible to the attentive probing pattern --- ngclearn/utils/analysis/attentive_probe.py | 39 ++++++++++++---------- 1 file changed, 22 insertions(+), 17 deletions(-) diff --git a/ngclearn/utils/analysis/attentive_probe.py b/ngclearn/utils/analysis/attentive_probe.py index e919eb49..482df014 100644 --- a/ngclearn/utils/analysis/attentive_probe.py +++ b/ngclearn/utils/analysis/attentive_probe.py @@ -102,11 +102,15 @@ def run_attention_probe(params, encodings, mask, n_heads: int, dropout: float = attn_params = (Wq, bq, Wk, bk, Wv, bv, Wout, bout) features = cross_attention(attn_params, learnable_query, encodings, mask, n_heads, dropout) features = features[:, 0] # (B, 1, dim) => (B, dim) - hids = jnp.matmul((features + learnable_query[:, 0]), Whid) + bhid - hids = gelu(hids) + # MLP + residual = features if use_LN: ## normalize hidden layer output of probe predictor - hids = layer_normalize(hids, Wln_mu, Wln_scale) - outs = jnp.matmul(hids, Wy) + by + features = layer_normalize(features, Wln_mu, Wln_scale) + features = jnp.matmul((features), Whid) + bhid + features = gelu(features) + features = residual + features + + outs = jnp.matmul(features, Wy) + by if use_softmax: ## apply softmax output nonlinearity outs = softmax(outs) return outs, features @@ -178,10 +182,11 @@ class AttentiveProbe(Probe): """ def __init__( - self, dkey, source_seq_length, input_dim, out_dim, num_heads=8, head_dim=64, + self, dkey, source_seq_length, input_dim, out_dim, num_heads=8, attn_dim=64, target_seq_length=1, learnable_query_dim=31, batch_size=1, hid_dim=32, use_LN=True, use_softmax=True, **kwargs ): super().__init__(dkey, batch_size, **kwargs) + assert attn_dim % num_heads == 0, f"`attn_dim` must be divisible by `num_heads`. Got {attn_dim} and {num_heads}." self.dkey, *subkeys = random.split(self.dkey, 12) self.num_heads = num_heads self.source_seq_length = source_seq_length @@ -192,24 +197,24 @@ def __init__( sigma = 0.05 ## cross-attention parameters - Wq = random.normal(subkeys[0], (learnable_query_dim, head_dim)) * sigma - bq = random.normal(subkeys[1], (1, head_dim)) * sigma - Wk = random.normal(subkeys[2], (input_dim, head_dim)) * sigma - bk = random.normal(subkeys[3], (1, head_dim)) * sigma - Wv = random.normal(subkeys[4], (input_dim, head_dim)) * sigma - bv = random.normal(subkeys[5], (1, head_dim)) * sigma - Wout = random.normal(subkeys[6], (head_dim, learnable_query_dim)) * sigma + Wq = random.normal(subkeys[0], (learnable_query_dim, attn_dim)) * sigma + bq = random.normal(subkeys[1], (1, attn_dim)) * sigma + Wk = random.normal(subkeys[2], (input_dim, attn_dim)) * sigma + bk = random.normal(subkeys[3], (1, attn_dim)) * sigma + Wv = random.normal(subkeys[4], (input_dim, attn_dim)) * sigma + bv = random.normal(subkeys[5], (1, attn_dim)) * sigma + Wout = random.normal(subkeys[6], (attn_dim, learnable_query_dim)) * sigma bout = random.normal(subkeys[7], (1, learnable_query_dim)) * sigma #params = (Wq, bq, Wk, bk, Wv, bv, Wout, bout) learnable_query = jnp.zeros((batch_size, 1, learnable_query_dim)) # (B, T, D) #self.all_params = (learnable_query, *params) self.mask = np.zeros((batch_size, target_seq_length, source_seq_length)).astype(bool) ## mask tensor ## MLP parameters - Whid = random.normal(subkeys[8], (learnable_query_dim, hid_dim)) * sigma - bhid = random.normal(subkeys[9], (1, hid_dim)) * sigma - Wln_mu = jnp.zeros((1, hid_dim)) - Wln_scale = jnp.ones((1, hid_dim)) - Wy = random.normal(subkeys[8], (hid_dim, out_dim)) * sigma + Whid = random.normal(subkeys[8], (learnable_query_dim, learnable_query_dim)) * sigma + bhid = random.normal(subkeys[9], (1, learnable_query_dim)) * sigma + Wln_mu = jnp.zeros((1, learnable_query_dim)) + Wln_scale = jnp.ones((1, learnable_query_dim)) + Wy = random.normal(subkeys[8], (learnable_query_dim, out_dim)) * sigma by = random.normal(subkeys[9], (1, out_dim)) * sigma #mlp_params = (Whid, bhid, Wln_mu, Wln_scale, Wy, by) self.probe_params = (learnable_query, Wq, bq, Wk, bk, Wv, bv, Wout, bout, Whid, bhid, Wln_mu, Wln_scale, Wy, by) From 8a36e407e841ac8f194f52f85dec32861e64ecf6 Mon Sep 17 00:00:00 2001 From: Viet Dung Nguyen Date: Sat, 1 Mar 2025 18:23:04 -0500 Subject: [PATCH 35/61] in layer normalization or any other Gaussian, standardeviation can never be zero. Additionally, if the subtraction inside the square root goes to zero, the gradient will become NaN. Therefore, adding a clipping is necessary. --- ngclearn/utils/model_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ngclearn/utils/model_utils.py b/ngclearn/utils/model_utils.py index 02adc01b..e5d8fcd5 100755 --- a/ngclearn/utils/model_utils.py +++ b/ngclearn/utils/model_utils.py @@ -548,7 +548,7 @@ def layer_normalize(x, shift=0., scale=1.): layer-normalized data samples `x` """ xmu = jnp.mean(x, axis=1, keepdims=True) - xsigma = jnp.sqrt(jnp.mean(jnp.square(x - xmu))) + xsigma = jnp.sqrt(jnp.mean(jnp.square(x - xmu)).clip(min=1e-6)) _x = (x - xmu)/(xsigma + 1e-6) return _x * scale + shift From f402d988736d9e13dae259b0fcdbdb4b1fd6d554 Mon Sep 17 00:00:00 2001 From: Viet Dung Nguyen Date: Sat, 1 Mar 2025 18:23:14 -0500 Subject: [PATCH 36/61] update attentive probe code --- ngclearn/utils/analysis/attentive_probe.py | 78 ++++++++++++++++------ 1 file changed, 57 insertions(+), 21 deletions(-) diff --git a/ngclearn/utils/analysis/attentive_probe.py b/ngclearn/utils/analysis/attentive_probe.py index 482df014..5c1707d4 100644 --- a/ngclearn/utils/analysis/attentive_probe.py +++ b/ngclearn/utils/analysis/attentive_probe.py @@ -78,7 +78,7 @@ def run_attention_probe(params, encodings, mask, n_heads: int, dropout: float = """ Runs full nonlinear attentive probe on input encodings (typically embedding vectors produced by some other model). - Args: + Args: params: parameters tuple/list of probe encodings: input encoding vectors/data @@ -98,18 +98,35 @@ def run_attention_probe(params, encodings, mask, n_heads: int, dropout: float = """ # encoded_image_feature: (B, hw, dim) #learnable_query, *_params) = params - learnable_query, Wq, bq, Wk, bk, Wv, bv, Wout, bout, Whid, bhid, Wln_mu, Wln_scale, Wy, by = params - attn_params = (Wq, bq, Wk, bk, Wv, bv, Wout, bout) - features = cross_attention(attn_params, learnable_query, encodings, mask, n_heads, dropout) + learnable_query, Wq, bq, Wk, bk, Wv, bv, Wout, bout,\ + Wqs, bqs, Wks, bks, Wvs, bvs, Wouts, bouts, Wlnattn_mu,\ + Wlnattn_scale, Whid1, bhid1, Wln_mu1, Wln_scale1, Whid2,\ + bhid2, Wln_mu2, Wln_scale2, Whid3, bhid3, Wln_mu3, Wln_scale3, Wy, by = params + cross_attn_params = (Wq, bq, Wk, bk, Wv, bv, Wout, bout) + features = cross_attention(cross_attn_params, learnable_query, encodings, mask, n_heads, dropout) + # Perform a single self-attention block here + # Self-Attention + self_attn_params = (Wqs, bqs, Wks, bks, Wvs, bvs, Wouts, bouts) + skip = features + if use_LN: + features = layer_normalize(features, Wlnattn_mu, Wlnattn_scale) + features = cross_attention(self_attn_params, features, features, None, n_heads, dropout) + features = features + skip features = features[:, 0] # (B, 1, dim) => (B, dim) # MLP - residual = features + skip = features if use_LN: ## normalize hidden layer output of probe predictor - features = layer_normalize(features, Wln_mu, Wln_scale) - features = jnp.matmul((features), Whid) + bhid + features = layer_normalize(features, Wln_mu1, Wln_scale1) + features = jnp.matmul((features), Whid1) + bhid1 features = gelu(features) - features = residual + features - + if use_LN: ## normalize hidden layer output of probe predictor + features = layer_normalize(features, Wln_mu2, Wln_scale2) + features = jnp.matmul((features), Whid2) + bhid2 + features = gelu(features) + if use_LN: ## normalize hidden layer output of probe predictor + features = layer_normalize(features, Wln_mu3, Wln_scale3) + features = jnp.matmul((features), Whid3) + bhid3 + features = features + skip outs = jnp.matmul(features, Wy) + by if use_softmax: ## apply softmax output nonlinearity outs = softmax(outs) @@ -183,11 +200,12 @@ class AttentiveProbe(Probe): """ def __init__( self, dkey, source_seq_length, input_dim, out_dim, num_heads=8, attn_dim=64, - target_seq_length=1, learnable_query_dim=31, batch_size=1, hid_dim=32, use_LN=True, use_softmax=True, **kwargs + target_seq_length=1, learnable_query_dim=32, batch_size=1, hid_dim=32, use_LN=True, use_softmax=True, **kwargs ): super().__init__(dkey, batch_size, **kwargs) assert attn_dim % num_heads == 0, f"`attn_dim` must be divisible by `num_heads`. Got {attn_dim} and {num_heads}." - self.dkey, *subkeys = random.split(self.dkey, 12) + assert learnable_query_dim % num_heads == 0, f"`learnable_query_dim` must be divisible by `num_heads`. Got {learnable_query_dim} and {num_heads}." + self.dkey, *subkeys = random.split(self.dkey, 25) self.num_heads = num_heads self.source_seq_length = source_seq_length self.input_dim = input_dim @@ -205,19 +223,37 @@ def __init__( bv = random.normal(subkeys[5], (1, attn_dim)) * sigma Wout = random.normal(subkeys[6], (attn_dim, learnable_query_dim)) * sigma bout = random.normal(subkeys[7], (1, learnable_query_dim)) * sigma - #params = (Wq, bq, Wk, bk, Wv, bv, Wout, bout) + cross_attn_params = (Wq, bq, Wk, bk, Wv, bv, Wout, bout) + Wqs = random.normal(subkeys[8], (learnable_query_dim, learnable_query_dim)) * sigma + bqs = random.normal(subkeys[9], (1, learnable_query_dim)) * sigma + Wks = random.normal(subkeys[10], (learnable_query_dim, learnable_query_dim)) * sigma + bks = random.normal(subkeys[11], (1, learnable_query_dim)) * sigma + Wvs = random.normal(subkeys[12], (learnable_query_dim, learnable_query_dim)) * sigma + bvs = random.normal(subkeys[13], (1, learnable_query_dim)) * sigma + Wouts = random.normal(subkeys[14], (learnable_query_dim, learnable_query_dim)) * sigma + bouts = random.normal(subkeys[15], (1, learnable_query_dim)) * sigma + Wlnattn_mu = jnp.zeros((1, learnable_query_dim)) + Wlnattn_scale = jnp.ones((1, learnable_query_dim)) + self_attn_params = (Wqs, bqs, Wks, bks, Wvs, bvs, Wouts, bouts, Wlnattn_mu, Wlnattn_scale) learnable_query = jnp.zeros((batch_size, 1, learnable_query_dim)) # (B, T, D) - #self.all_params = (learnable_query, *params) self.mask = np.zeros((batch_size, target_seq_length, source_seq_length)).astype(bool) ## mask tensor ## MLP parameters - Whid = random.normal(subkeys[8], (learnable_query_dim, learnable_query_dim)) * sigma - bhid = random.normal(subkeys[9], (1, learnable_query_dim)) * sigma - Wln_mu = jnp.zeros((1, learnable_query_dim)) - Wln_scale = jnp.ones((1, learnable_query_dim)) - Wy = random.normal(subkeys[8], (learnable_query_dim, out_dim)) * sigma - by = random.normal(subkeys[9], (1, out_dim)) * sigma - #mlp_params = (Whid, bhid, Wln_mu, Wln_scale, Wy, by) - self.probe_params = (learnable_query, Wq, bq, Wk, bk, Wv, bv, Wout, bout, Whid, bhid, Wln_mu, Wln_scale, Wy, by) + Whid1 = random.normal(subkeys[16], (learnable_query_dim, learnable_query_dim)) * sigma + bhid1 = random.normal(subkeys[17], (1, learnable_query_dim)) * sigma + Wln_mu1 = jnp.zeros((1, learnable_query_dim)) + Wln_scale1 = jnp.ones((1, learnable_query_dim)) + Whid2 = random.normal(subkeys[18], (learnable_query_dim, learnable_query_dim * 4)) * sigma + bhid2 = random.normal(subkeys[19], (1, learnable_query_dim * 4)) * sigma + Wln_mu2 = jnp.zeros((1, learnable_query_dim)) + Wln_scale2 = jnp.ones((1, learnable_query_dim)) + Whid3 = random.normal(subkeys[20], (learnable_query_dim * 4, learnable_query_dim)) * sigma + bhid3 = random.normal(subkeys[21], (1, learnable_query_dim)) * sigma + Wln_mu3 = jnp.zeros((1, learnable_query_dim * 4)) + Wln_scale3 = jnp.ones((1, learnable_query_dim * 4)) + Wy = random.normal(subkeys[22], (learnable_query_dim, out_dim)) * sigma + by = random.normal(subkeys[23], (1, out_dim)) * sigma + mlp_params = (Whid1, bhid1, Wln_mu1, Wln_scale1, Whid2, bhid2, Wln_mu2, Wln_scale2, Whid3, bhid3, Wln_mu3, Wln_scale3, Wy, by) + self.probe_params = (learnable_query, *cross_attn_params, *self_attn_params, *mlp_params) ## set up gradient calculator self.grad_fx = jax.value_and_grad(eval_attention_probe, argnums=0, has_aux=True) From 2a71b7f516c933738245f88ce02bbc74500b37d0 Mon Sep 17 00:00:00 2001 From: Alexander Ororbia Date: Mon, 3 Mar 2025 11:41:19 -0500 Subject: [PATCH 37/61] minor tweak to attentive prob code comments --- ngclearn/utils/analysis/attentive_probe.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/ngclearn/utils/analysis/attentive_probe.py b/ngclearn/utils/analysis/attentive_probe.py index 5c1707d4..751f101e 100644 --- a/ngclearn/utils/analysis/attentive_probe.py +++ b/ngclearn/utils/analysis/attentive_probe.py @@ -232,24 +232,24 @@ def __init__( bvs = random.normal(subkeys[13], (1, learnable_query_dim)) * sigma Wouts = random.normal(subkeys[14], (learnable_query_dim, learnable_query_dim)) * sigma bouts = random.normal(subkeys[15], (1, learnable_query_dim)) * sigma - Wlnattn_mu = jnp.zeros((1, learnable_query_dim)) - Wlnattn_scale = jnp.ones((1, learnable_query_dim)) + Wlnattn_mu = jnp.zeros((1, learnable_query_dim)) ## LN parameter (applied to output of attention) + Wlnattn_scale = jnp.ones((1, learnable_query_dim)) ## LN parameter (applied to output of attention) self_attn_params = (Wqs, bqs, Wks, bks, Wvs, bvs, Wouts, bouts, Wlnattn_mu, Wlnattn_scale) learnable_query = jnp.zeros((batch_size, 1, learnable_query_dim)) # (B, T, D) self.mask = np.zeros((batch_size, target_seq_length, source_seq_length)).astype(bool) ## mask tensor ## MLP parameters Whid1 = random.normal(subkeys[16], (learnable_query_dim, learnable_query_dim)) * sigma bhid1 = random.normal(subkeys[17], (1, learnable_query_dim)) * sigma - Wln_mu1 = jnp.zeros((1, learnable_query_dim)) - Wln_scale1 = jnp.ones((1, learnable_query_dim)) + Wln_mu1 = jnp.zeros((1, learnable_query_dim)) ## LN parameter + Wln_scale1 = jnp.ones((1, learnable_query_dim)) ## LN parameter Whid2 = random.normal(subkeys[18], (learnable_query_dim, learnable_query_dim * 4)) * sigma bhid2 = random.normal(subkeys[19], (1, learnable_query_dim * 4)) * sigma - Wln_mu2 = jnp.zeros((1, learnable_query_dim)) - Wln_scale2 = jnp.ones((1, learnable_query_dim)) + Wln_mu2 = jnp.zeros((1, learnable_query_dim)) ## LN parameter + Wln_scale2 = jnp.ones((1, learnable_query_dim)) ## LN parameter Whid3 = random.normal(subkeys[20], (learnable_query_dim * 4, learnable_query_dim)) * sigma bhid3 = random.normal(subkeys[21], (1, learnable_query_dim)) * sigma - Wln_mu3 = jnp.zeros((1, learnable_query_dim * 4)) - Wln_scale3 = jnp.ones((1, learnable_query_dim * 4)) + Wln_mu3 = jnp.zeros((1, learnable_query_dim * 4)) ## LN parameter + Wln_scale3 = jnp.ones((1, learnable_query_dim * 4)) ## LN parameter Wy = random.normal(subkeys[22], (learnable_query_dim, out_dim)) * sigma by = random.normal(subkeys[23], (1, out_dim)) * sigma mlp_params = (Whid1, bhid1, Wln_mu1, Wln_scale1, Whid2, bhid2, Wln_mu2, Wln_scale2, Whid3, bhid3, Wln_mu3, Wln_scale3, Wy, by) From b688c6c401c453f4f52961190d29f06405559456 Mon Sep 17 00:00:00 2001 From: Alexander Ororbia Date: Mon, 3 Mar 2025 18:12:43 -0500 Subject: [PATCH 38/61] cleaned up probe parent fit routine --- ngclearn/utils/analysis/probe.py | 60 +++++++++++++++++++++++++------- 1 file changed, 48 insertions(+), 12 deletions(-) diff --git a/ngclearn/utils/analysis/probe.py b/ngclearn/utils/analysis/probe.py index d9aa1cf2..84100fe1 100644 --- a/ngclearn/utils/analysis/probe.py +++ b/ngclearn/utils/analysis/probe.py @@ -53,18 +53,29 @@ def predict(self, data): Y_mu = jnp.concatenate(Y_mu, axis=0) return Y_mu - def fit(self, data, labels, n_iter=50): + def fit(self, dataset, dev_dataset=None, n_iter=50, patience=20): """ Fits this probe to a pool of data. Args: - data: a dataset or design tensor/matrix containing encoding vector sequences; shape (N, T, embed_dim) or (N, embed_dim) + dataset: a dataset tuple containing two design tensors/matrices (X, Y), with the first containing encoding + vector sequences of shape (N, T, embed_dim) or (N, embed_dim) and the second containing the + corresponding labels/targets for the embedding data of shape (N, target_dim); (Default: None) + + dev_dataset: an optional development set tuple, with same format as `dataset` (Default: None) + + n_iter: number of iterations to run model fitting (Default: 50 iterations) - labels: a design matrix containing corresponding labels/targets for the embedding data; shape (N, target_dim) + patience: number of iterations of improvement (decrease) in loss before early-stopping enacted Returns: the output scores/predictions made by this probe """ + data, labels = dataset + dev_data = dev_labels = None + if dev_dataset is not None: + dev_data, dev_labels = dev_dataset + _data = data if len(_data.shape) < 3: _data = jnp.expand_dims(_data, axis=1) @@ -79,9 +90,12 @@ def fit(self, data, labels, n_iter=50): n_batches = int(n_samples / self.batch_size) ## run main probe fitting loop - Y_mu = [] + impatience = 0 + final_L = 10000. + best_acc = 0. + #Y_mu = [] _Y = None - for iter in range(n_iter): + for ii in range(n_iter): ## shuffle data (to ensure i.i.d. across sequences) self.dkey, *subkeys = random.split(self.dkey, 2) ptrs = random.permutation(subkeys[0], n_samples) @@ -89,6 +103,7 @@ def fit(self, data, labels, n_iter=50): _Y = labels[ptrs, :] ## run one epoch over data tensors L = 0. + acc = 0. Ns = 0. s_ptr = 0 @@ -101,12 +116,33 @@ def fit(self, data, labels, n_iter=50): Ns += x_mb.shape[0] _L, py = self.update(x_mb, y_mb) - L = _L + L - print(f"\r{iter} L = {L/Ns}", end="") # p(y|z):\n{py}") - if iter == n_iter-1: - Y_mu.append(py) + acc = jnp.sum(jnp.equal(jnp.argmax(py, axis=1), jnp.argmax(y_mb, axis=1))) + acc + L = (_L * x_mb.shape[0]) + L ## we remove the batch division from loss w.r.t. x_mb/y_mb + if dev_data is not None: + print(f"\r{ii} L = {L / Ns:.3f} Acc = {acc / Ns:.2f} Dev.Acc = {best_acc:.2f}", end="") + else: + print(f"\r{ii} L = {L / Ns:.3f} Acc = {acc / Ns:.2f}", end="") + # if ii == ii-1: + # Y_mu.append(py) print() - if iter == n_iter - 1: - Y_mu = jnp.concatenate(Y_mu, axis=0) - return Y_mu, _Y ## return predictions mapped to current shuffling of labels + acc = acc / Ns + final_L = L / Ns ## compute current loss over (train) dataset + # if ii == ii - 1: + # Y_mu = jnp.concatenate(Y_mu, axis=0) + + impatience += 1 + if dev_data is not None: + Ymu = self.predict(dev_data) + acc = jnp.sum(jnp.equal(jnp.argmax(Ymu, axis=1), jnp.argmax(dev_labels, axis=1))) / (dev_labels.shape[0] * 1.) + if acc > best_acc: + best_acc = acc + impatience = 0 + else: ## use training acc if no dev-set provided + if acc > best_acc: + best_acc = acc + impatience = 0 + + if impatience > patience: + break ## execute early stopping + return final_L From 9ad4ae29f716f5fe0afb77f18aa2a932214ee922 Mon Sep 17 00:00:00 2001 From: Alexander Ororbia Date: Mon, 3 Mar 2025 18:21:41 -0500 Subject: [PATCH 39/61] cleaned up probe parent fit routine --- ngclearn/utils/analysis/linear_probe.py | 1 + ngclearn/utils/analysis/probe.py | 16 +++++++++++----- 2 files changed, 12 insertions(+), 5 deletions(-) diff --git a/ngclearn/utils/analysis/linear_probe.py b/ngclearn/utils/analysis/linear_probe.py index 05284e94..370000f4 100644 --- a/ngclearn/utils/analysis/linear_probe.py +++ b/ngclearn/utils/analysis/linear_probe.py @@ -84,6 +84,7 @@ def __init__( self.use_LN = use_LN self.l2_decay = 0.0001 self.l1_decay = 0.000025 + # eta = 0.05 for SGD, batch_size=2000 ## set up classifier flat_input_dim = input_dim * source_seq_length diff --git a/ngclearn/utils/analysis/probe.py b/ngclearn/utils/analysis/probe.py index 84100fe1..74bb16c5 100644 --- a/ngclearn/utils/analysis/probe.py +++ b/ngclearn/utils/analysis/probe.py @@ -11,11 +11,12 @@ class Probe(): """ def __init__( - self, dkey, batch_size=4, **kwargs + self, dkey, batch_size=1, dev_batch_size=1, **kwargs ): #dkey, *subkeys = random.split(dkey, 3) self.dkey = dkey self.batch_size = batch_size + self.dev_batch_size = dev_batch_size def process(self, embeddings): predictions = None @@ -25,24 +26,29 @@ def update(self, embeddings, labels): L = predictions = None return L, predictions - def predict(self, data): + def predict(self, data, batch_size=None): """ Runs this probe's inference scheme over a pool of data. Args: data: a dataset or design tensor/matrix containing encoding vector sequences; shape (N, T, embed_dim) or (N, embed_dim) + batch_size: optional batch-size argument (Default: None, will use training batch size) + Returns: the output scores/predictions made by this probe """ + _batch_size = batch_size + if _batch_size is None: + _batch_size = self.batch_size _data = data if len(_data.shape) < 3: _data = jnp.expand_dims(_data, axis=1) n_samples, seq_len, dim = _data.shape - n_batches = int(n_samples / self.batch_size) + n_batches = int(n_samples / _batch_size) s_ptr = 0 - e_ptr = self.batch_size + e_ptr = _batch_size Y_mu = [] for b in range(n_batches): x_mb = _data[s_ptr:e_ptr, :, :] ## slice out 3D batch tensor @@ -132,7 +138,7 @@ def fit(self, dataset, dev_dataset=None, n_iter=50, patience=20): impatience += 1 if dev_data is not None: - Ymu = self.predict(dev_data) + Ymu = self.predict(dev_data, batch_size=self.dev_batch_size) acc = jnp.sum(jnp.equal(jnp.argmax(Ymu, axis=1), jnp.argmax(dev_labels, axis=1))) / (dev_labels.shape[0] * 1.) if acc > best_acc: best_acc = acc From 3a2de992f723fbb26c1e635217377f240bc9350b Mon Sep 17 00:00:00 2001 From: Alexander Ororbia Date: Mon, 3 Mar 2025 18:28:50 -0500 Subject: [PATCH 40/61] cleaned up probe parent fit routine --- ngclearn/utils/analysis/probe.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/ngclearn/utils/analysis/probe.py b/ngclearn/utils/analysis/probe.py index 74bb16c5..ac5a304e 100644 --- a/ngclearn/utils/analysis/probe.py +++ b/ngclearn/utils/analysis/probe.py @@ -75,7 +75,7 @@ def fit(self, dataset, dev_dataset=None, n_iter=50, patience=20): patience: number of iterations of improvement (decrease) in loss before early-stopping enacted Returns: - the output scores/predictions made by this probe + best accuracy found over fitting run """ data, labels = dataset dev_data = dev_labels = None @@ -97,9 +97,7 @@ def fit(self, dataset, dev_dataset=None, n_iter=50, patience=20): ## run main probe fitting loop impatience = 0 - final_L = 10000. best_acc = 0. - #Y_mu = [] _Y = None for ii in range(n_iter): ## shuffle data (to ensure i.i.d. across sequences) @@ -128,13 +126,9 @@ def fit(self, dataset, dev_dataset=None, n_iter=50, patience=20): print(f"\r{ii} L = {L / Ns:.3f} Acc = {acc / Ns:.2f} Dev.Acc = {best_acc:.2f}", end="") else: print(f"\r{ii} L = {L / Ns:.3f} Acc = {acc / Ns:.2f}", end="") - # if ii == ii-1: - # Y_mu.append(py) print() acc = acc / Ns - final_L = L / Ns ## compute current loss over (train) dataset - # if ii == ii - 1: - # Y_mu = jnp.concatenate(Y_mu, axis=0) + L = L / Ns ## compute current loss over (train) dataset impatience += 1 if dev_data is not None: @@ -150,5 +144,5 @@ def fit(self, dataset, dev_dataset=None, n_iter=50, patience=20): if impatience > patience: break ## execute early stopping - return final_L + return best_acc From 155d8301d0a6bcf4d013621d291d41b5c9f8e797 Mon Sep 17 00:00:00 2001 From: Alexander Ororbia Date: Mon, 3 Mar 2025 18:49:22 -0500 Subject: [PATCH 41/61] cleaned up probe parent fit routine --- ngclearn/utils/analysis/probe.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/ngclearn/utils/analysis/probe.py b/ngclearn/utils/analysis/probe.py index ac5a304e..7ad13238 100644 --- a/ngclearn/utils/analysis/probe.py +++ b/ngclearn/utils/analysis/probe.py @@ -122,11 +122,12 @@ def fit(self, dataset, dev_dataset=None, n_iter=50, patience=20): _L, py = self.update(x_mb, y_mb) acc = jnp.sum(jnp.equal(jnp.argmax(py, axis=1), jnp.argmax(y_mb, axis=1))) + acc L = (_L * x_mb.shape[0]) + L ## we remove the batch division from loss w.r.t. x_mb/y_mb - if dev_data is not None: - print(f"\r{ii} L = {L / Ns:.3f} Acc = {acc / Ns:.2f} Dev.Acc = {best_acc:.2f}", end="") - else: - print(f"\r{ii} L = {L / Ns:.3f} Acc = {acc / Ns:.2f}", end="") - print() + + if dev_data is not None: + print(f"\r{ii} L = {L / Ns:.3f} Acc = {acc / Ns:.2f} Dev.Acc = {best_acc:.2f}", end="") + else: + print(f"\r{ii} L = {L / Ns:.3f} Acc = {acc / Ns:.2f}", end="") + acc = acc / Ns L = L / Ns ## compute current loss over (train) dataset @@ -144,5 +145,6 @@ def fit(self, dataset, dev_dataset=None, n_iter=50, patience=20): if impatience > patience: break ## execute early stopping + print() return best_acc From 099c588756e5b6e650ed5caba4e03b69ed0d30a9 Mon Sep 17 00:00:00 2001 From: Alexander Ororbia Date: Wed, 5 Mar 2025 13:55:10 -0500 Subject: [PATCH 42/61] minor edits to attn probe --- ngclearn/utils/analysis/attentive_probe.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/ngclearn/utils/analysis/attentive_probe.py b/ngclearn/utils/analysis/attentive_probe.py index 751f101e..aa25f51b 100644 --- a/ngclearn/utils/analysis/attentive_probe.py +++ b/ngclearn/utils/analysis/attentive_probe.py @@ -236,7 +236,8 @@ def __init__( Wlnattn_scale = jnp.ones((1, learnable_query_dim)) ## LN parameter (applied to output of attention) self_attn_params = (Wqs, bqs, Wks, bks, Wvs, bvs, Wouts, bouts, Wlnattn_mu, Wlnattn_scale) learnable_query = jnp.zeros((batch_size, 1, learnable_query_dim)) # (B, T, D) - self.mask = np.zeros((batch_size, target_seq_length, source_seq_length)).astype(bool) ## mask tensor + self.mask = np.zeros((self.batch_size, target_seq_length, source_seq_length)).astype(bool) ## mask tensor + self.dev_mask = np.zeros((self.dev_batch_size, target_seq_length, source_seq_length)).astype(bool) ## MLP parameters Whid1 = random.normal(subkeys[16], (learnable_query_dim, learnable_query_dim)) * sigma bhid1 = random.normal(subkeys[17], (1, learnable_query_dim)) * sigma @@ -259,7 +260,7 @@ def __init__( self.grad_fx = jax.value_and_grad(eval_attention_probe, argnums=0, has_aux=True) ## set up update rule/optimizer self.optim_params = adam.adam_init(self.probe_params) - self.eta = 0.001 + self.eta = 0.0002 #0.001 def process(self, embedding_sequence): """ @@ -271,13 +272,14 @@ def process(self, embedding_sequence): Returns: probe output scores/probability values """ + #print(embedding_sequence.shape) outs, feats = run_attention_probe( - self.probe_params, embedding_sequence, self.mask, self.num_heads, 0.0, use_LN=self.use_LN, + self.probe_params, embedding_sequence, self.dev_mask, self.num_heads, 0.0, use_LN=self.use_LN, use_softmax=self.use_softmax ) return outs - def update(self, embedding_sequence, labels): + def update(self, embedding_sequence, labels, dkey=None): """ Runs and updates this probe given an input batch of sequences of encodings/embeddings and their externally assigned labels/target vector values. @@ -290,9 +292,10 @@ def update(self, embedding_sequence, labels): Returns: probe output scores/probability values """ + # TODO: put in dkey to facilitate dropout ## compute partial derivatives / adjustments to probe parameters outputs, grads = self.grad_fx( - self.probe_params, embedding_sequence, labels, self.mask, self.num_heads, dropout=0., use_LN=self.use_LN, + self.probe_params, embedding_sequence, labels, self.mask, self.num_heads, dropout=0.5, use_LN=self.use_LN, use_softmax=self.use_softmax ) loss, predictions = outputs From aeabf61285962a208e38d5f387b560a8bdbcd512 Mon Sep 17 00:00:00 2001 From: Viet Nguyen Date: Wed, 5 Mar 2025 18:53:34 -0500 Subject: [PATCH 43/61] update attentive probe with input layer norm --- ngclearn/utils/analysis/attentive_probe.py | 21 +++++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/ngclearn/utils/analysis/attentive_probe.py b/ngclearn/utils/analysis/attentive_probe.py index aa25f51b..de93a2f3 100644 --- a/ngclearn/utils/analysis/attentive_probe.py +++ b/ngclearn/utils/analysis/attentive_probe.py @@ -73,8 +73,8 @@ def cross_attention(params: tuple, x1: jax.Array, x2: jax.Array, mask: jax.Array attention = attention.transpose([0, 2, 1, 3]).reshape((B, T, -1)) # (B, T, H, E) => (B, T, D) return attention @ Wout + bout # (B, T, Dq) -@bind(jax.jit, static_argnums=[3, 4, 5, 6]) -def run_attention_probe(params, encodings, mask, n_heads: int, dropout: float = 0.0, use_LN=False, use_softmax=True): +@bind(jax.jit, static_argnums=[3, 4, 5, 6, 7]) +def run_attention_probe(params, encodings, mask, n_heads: int, dropout: float = 0.0, use_LN=False, use_LN_input=True, use_softmax=True): """ Runs full nonlinear attentive probe on input encodings (typically embedding vectors produced by some other model). @@ -101,8 +101,11 @@ def run_attention_probe(params, encodings, mask, n_heads: int, dropout: float = learnable_query, Wq, bq, Wk, bk, Wv, bv, Wout, bout,\ Wqs, bqs, Wks, bks, Wvs, bvs, Wouts, bouts, Wlnattn_mu,\ Wlnattn_scale, Whid1, bhid1, Wln_mu1, Wln_scale1, Whid2,\ - bhid2, Wln_mu2, Wln_scale2, Whid3, bhid3, Wln_mu3, Wln_scale3, Wy, by = params + bhid2, Wln_mu2, Wln_scale2, Whid3, bhid3, Wln_mu3, Wln_scale3,\ + Wy, by, ln_in_mu, ln_in_scale = params cross_attn_params = (Wq, bq, Wk, bk, Wv, bv, Wout, bout) + if use_LN_input: + learnable_query = layer_normalize(learnable_query, ln_in_mu, ln_in_scale) features = cross_attention(cross_attn_params, learnable_query, encodings, mask, n_heads, dropout) # Perform a single self-attention block here # Self-Attention @@ -200,7 +203,7 @@ class AttentiveProbe(Probe): """ def __init__( self, dkey, source_seq_length, input_dim, out_dim, num_heads=8, attn_dim=64, - target_seq_length=1, learnable_query_dim=32, batch_size=1, hid_dim=32, use_LN=True, use_softmax=True, **kwargs + target_seq_length=1, learnable_query_dim=32, batch_size=1, hid_dim=32, use_LN=True, use_LN_input=True, use_softmax=True, **kwargs ): super().__init__(dkey, batch_size, **kwargs) assert attn_dim % num_heads == 0, f"`attn_dim` must be divisible by `num_heads`. Got {attn_dim} and {num_heads}." @@ -212,6 +215,7 @@ def __init__( self.out_dim = out_dim self.use_softmax = use_softmax self.use_LN = use_LN + self.use_LN_input = use_LN_input sigma = 0.05 ## cross-attention parameters @@ -254,7 +258,11 @@ def __init__( Wy = random.normal(subkeys[22], (learnable_query_dim, out_dim)) * sigma by = random.normal(subkeys[23], (1, out_dim)) * sigma mlp_params = (Whid1, bhid1, Wln_mu1, Wln_scale1, Whid2, bhid2, Wln_mu2, Wln_scale2, Whid3, bhid3, Wln_mu3, Wln_scale3, Wy, by) - self.probe_params = (learnable_query, *cross_attn_params, *self_attn_params, *mlp_params) + # Finally, define ln for the input to the attention + ln_in_mu = jnp.zeros((1, learnable_query_dim)) ## LN parameter + ln_in_scale = jnp.ones((1, learnable_query_dim)) ## LN parameter + ln_in_params = (ln_in_mu, ln_in_scale) + self.probe_params = (learnable_query, *cross_attn_params, *self_attn_params, *mlp_params, *ln_in_params) ## set up gradient calculator self.grad_fx = jax.value_and_grad(eval_attention_probe, argnums=0, has_aux=True) @@ -294,8 +302,9 @@ def update(self, embedding_sequence, labels, dkey=None): """ # TODO: put in dkey to facilitate dropout ## compute partial derivatives / adjustments to probe parameters + # NOTE: Viet: Change back to 0.0 for now for the code to run outputs, grads = self.grad_fx( - self.probe_params, embedding_sequence, labels, self.mask, self.num_heads, dropout=0.5, use_LN=self.use_LN, + self.probe_params, embedding_sequence, labels, self.mask, self.num_heads, dropout=0.0, use_LN=self.use_LN, use_softmax=self.use_softmax ) loss, predictions = outputs From 8682954d13ec0c1f914759afe3285a00caf02f1c Mon Sep 17 00:00:00 2001 From: Viet Nguyen Date: Wed, 5 Mar 2025 19:18:48 -0500 Subject: [PATCH 44/61] update input layer normalization --- ngclearn/utils/analysis/attentive_probe.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/ngclearn/utils/analysis/attentive_probe.py b/ngclearn/utils/analysis/attentive_probe.py index de93a2f3..a5db0d12 100644 --- a/ngclearn/utils/analysis/attentive_probe.py +++ b/ngclearn/utils/analysis/attentive_probe.py @@ -102,10 +102,11 @@ def run_attention_probe(params, encodings, mask, n_heads: int, dropout: float = Wqs, bqs, Wks, bks, Wvs, bvs, Wouts, bouts, Wlnattn_mu,\ Wlnattn_scale, Whid1, bhid1, Wln_mu1, Wln_scale1, Whid2,\ bhid2, Wln_mu2, Wln_scale2, Whid3, bhid3, Wln_mu3, Wln_scale3,\ - Wy, by, ln_in_mu, ln_in_scale = params + Wy, by, ln_in_mu, ln_in_scale, ln_in_mu2, ln_in_scale2 = params cross_attn_params = (Wq, bq, Wk, bk, Wv, bv, Wout, bout) if use_LN_input: learnable_query = layer_normalize(learnable_query, ln_in_mu, ln_in_scale) + encodings = layer_normalize(encodings, ln_in_mu2, ln_in_scale2) features = cross_attention(cross_attn_params, learnable_query, encodings, mask, n_heads, dropout) # Perform a single self-attention block here # Self-Attention @@ -261,7 +262,9 @@ def __init__( # Finally, define ln for the input to the attention ln_in_mu = jnp.zeros((1, learnable_query_dim)) ## LN parameter ln_in_scale = jnp.ones((1, learnable_query_dim)) ## LN parameter - ln_in_params = (ln_in_mu, ln_in_scale) + ln_in_mu2 = jnp.zeros((1, input_dim)) ## LN parameter + ln_in_scale2 = jnp.ones((1, input_dim)) ## LN parameter + ln_in_params = (ln_in_mu, ln_in_scale, ln_in_mu2, ln_in_scale2) self.probe_params = (learnable_query, *cross_attn_params, *self_attn_params, *mlp_params, *ln_in_params) ## set up gradient calculator From dc8c12709f6f054eb3811105ac2a43470d14a084 Mon Sep 17 00:00:00 2001 From: Viet Nguyen Date: Wed, 5 Mar 2025 19:41:18 -0500 Subject: [PATCH 45/61] update code to fix nan bug --- ngclearn/utils/analysis/attentive_probe.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/ngclearn/utils/analysis/attentive_probe.py b/ngclearn/utils/analysis/attentive_probe.py index a5db0d12..c058d3c2 100644 --- a/ngclearn/utils/analysis/attentive_probe.py +++ b/ngclearn/utils/analysis/attentive_probe.py @@ -133,7 +133,9 @@ def run_attention_probe(params, encodings, mask, n_heads: int, dropout: float = features = features + skip outs = jnp.matmul(features, Wy) + by if use_softmax: ## apply softmax output nonlinearity - outs = softmax(outs) + # NOTE: Viet: please check the softmax function, it might potentially + # cause the gradient to be nan since there is a potential division by zero + outs = jax.nn.softmax(outs) return outs, features @bind(jax.jit, static_argnums=[4, 5, 6, 7]) @@ -165,7 +167,7 @@ def eval_attention_probe(params, encodings, labels, mask, n_heads: int, dropout: # encodings: (B, hw, dim) outs, _ = run_attention_probe(params, encodings, mask, n_heads, dropout, use_LN, use_softmax) if use_softmax: ## Multinoulli log likelihood for 1-of-K predictions - L = -jnp.mean(jnp.sum(jnp.log(outs) * labels, axis=1, keepdims=True)) + L = -jnp.mean(jnp.sum(jnp.log(outs.clip(min=1e-5)) * labels, axis=1, keepdims=True)) else: ## MSE for real-valued outputs L = jnp.mean(jnp.sum(jnp.square(outs - labels), axis=1, keepdims=True)) return L, outs #, features From 27fd9bfa71d4ace71c7712e775c45d12345d74c8 Mon Sep 17 00:00:00 2001 From: Alexander Ororbia Date: Wed, 5 Mar 2025 19:59:35 -0500 Subject: [PATCH 46/61] minor tweak to attn probe --- ngclearn/utils/analysis/attentive_probe.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/ngclearn/utils/analysis/attentive_probe.py b/ngclearn/utils/analysis/attentive_probe.py index c058d3c2..2a325c8d 100644 --- a/ngclearn/utils/analysis/attentive_probe.py +++ b/ngclearn/utils/analysis/attentive_probe.py @@ -74,7 +74,7 @@ def cross_attention(params: tuple, x1: jax.Array, x2: jax.Array, mask: jax.Array return attention @ Wout + bout # (B, T, Dq) @bind(jax.jit, static_argnums=[3, 4, 5, 6, 7]) -def run_attention_probe(params, encodings, mask, n_heads: int, dropout: float = 0.0, use_LN=False, use_LN_input=True, use_softmax=True): +def run_attention_probe(params, encodings, mask, n_heads: int, dropout: float = 0.0, use_LN=False, use_LN_input=False, use_softmax=True): """ Runs full nonlinear attentive probe on input encodings (typically embedding vectors produced by some other model). @@ -138,8 +138,8 @@ def run_attention_probe(params, encodings, mask, n_heads: int, dropout: float = outs = jax.nn.softmax(outs) return outs, features -@bind(jax.jit, static_argnums=[4, 5, 6, 7]) -def eval_attention_probe(params, encodings, labels, mask, n_heads: int, dropout: float = 0.0, use_LN=False, use_softmax=True): +@bind(jax.jit, static_argnums=[4, 5, 6, 7, 8]) +def eval_attention_probe(params, encodings, labels, mask, n_heads: int, dropout: float = 0.0, use_LN=False, use_LN_input=False, use_softmax=True): """ Runs and evaluates the nonlinear attentive probe given a paired set of encoding vectors and externally assigned labels/regression targets. @@ -165,7 +165,7 @@ def eval_attention_probe(params, encodings, labels, mask, n_heads: int, dropout: current loss value, output scores/probabilities """ # encodings: (B, hw, dim) - outs, _ = run_attention_probe(params, encodings, mask, n_heads, dropout, use_LN, use_softmax) + outs, _ = run_attention_probe(params, encodings, mask, n_heads, dropout, use_LN, use_LN_input, use_softmax) if use_softmax: ## Multinoulli log likelihood for 1-of-K predictions L = -jnp.mean(jnp.sum(jnp.log(outs.clip(min=1e-5)) * labels, axis=1, keepdims=True)) else: ## MSE for real-valued outputs @@ -206,7 +206,7 @@ class AttentiveProbe(Probe): """ def __init__( self, dkey, source_seq_length, input_dim, out_dim, num_heads=8, attn_dim=64, - target_seq_length=1, learnable_query_dim=32, batch_size=1, hid_dim=32, use_LN=True, use_LN_input=True, use_softmax=True, **kwargs + target_seq_length=1, learnable_query_dim=32, batch_size=1, hid_dim=32, use_LN=True, use_LN_input=False, use_softmax=True, **kwargs ): super().__init__(dkey, batch_size, **kwargs) assert attn_dim % num_heads == 0, f"`attn_dim` must be divisible by `num_heads`. Got {attn_dim} and {num_heads}." @@ -288,7 +288,7 @@ def process(self, embedding_sequence): #print(embedding_sequence.shape) outs, feats = run_attention_probe( self.probe_params, embedding_sequence, self.dev_mask, self.num_heads, 0.0, use_LN=self.use_LN, - use_softmax=self.use_softmax + use_LN_input=self.use_LN_input, use_softmax=self.use_softmax ) return outs @@ -310,7 +310,7 @@ def update(self, embedding_sequence, labels, dkey=None): # NOTE: Viet: Change back to 0.0 for now for the code to run outputs, grads = self.grad_fx( self.probe_params, embedding_sequence, labels, self.mask, self.num_heads, dropout=0.0, use_LN=self.use_LN, - use_softmax=self.use_softmax + use_LN_input=self.use_LN_input, use_softmax=self.use_softmax ) loss, predictions = outputs ## adjust parameters of probe From 84005b5685673abc7309fc0b7e3eba58964749c2 Mon Sep 17 00:00:00 2001 From: Alexander Ororbia Date: Wed, 5 Mar 2025 20:13:27 -0500 Subject: [PATCH 47/61] cleaned up probes --- ngclearn/utils/analysis/attentive_probe.py | 75 ++++++++++------------ ngclearn/utils/analysis/linear_probe.py | 25 +------- ngclearn/utils/analysis/probe.py | 29 ++++++++- 3 files changed, 64 insertions(+), 65 deletions(-) diff --git a/ngclearn/utils/analysis/attentive_probe.py b/ngclearn/utils/analysis/attentive_probe.py index 2a325c8d..7e17e953 100644 --- a/ngclearn/utils/analysis/attentive_probe.py +++ b/ngclearn/utils/analysis/attentive_probe.py @@ -21,8 +21,8 @@ def masked_fill(x: jax.Array, mask: jax.Array, value=0) -> jax.Array: """ return jnp.where(mask, jnp.broadcast_to(value, x.shape), x) -@bind(jax.jit, static_argnums=[4, 5]) -def cross_attention(params: tuple, x1: jax.Array, x2: jax.Array, mask: jax.Array, n_heads: int=8, dropout_rate: float=0.0) -> jax.Array: +@bind(jax.jit, static_argnums=[5, 6]) +def cross_attention(dkey, params: tuple, x1: jax.Array, x2: jax.Array, mask: jax.Array, n_heads: int=8, dropout_rate: float=0.0) -> jax.Array: """ Run cross-attention function given a list of parameters and two sequences (x1 and x2). The function takes in a query sequence x1 and a key-value sequence x2, and returns an output of the same shape as x1. @@ -31,6 +31,8 @@ def cross_attention(params: tuple, x1: jax.Array, x2: jax.Array, mask: jax.Array H is the number of attention heads. Args: + dkey: JAX key to trigger any internal noise (drop-out) + params (tuple): tuple of parameters x1 (jax.Array): query sequence. Shape: (B, T, Dq) @@ -68,17 +70,22 @@ def cross_attention(params: tuple, x1: jax.Array, x2: jax.Array, mask: jax.Array score = jax.nn.softmax(score, axis=-1) # (B, H, T, S) score = score.astype(q.dtype) # (B, H, T, S) if dropout_rate > 0.: - score = drop_out(input=score, rate=dropout_rate) ## NOTE: normally you apply dropout here + score = drop_out(dkey, input=score, rate=dropout_rate) ## NOTE: normally you apply dropout here attention = jnp.einsum("BHTS,BHSE->BHTE", score, v) # (B, T, H, E) attention = attention.transpose([0, 2, 1, 3]).reshape((B, T, -1)) # (B, T, H, E) => (B, T, D) return attention @ Wout + bout # (B, T, Dq) -@bind(jax.jit, static_argnums=[3, 4, 5, 6, 7]) -def run_attention_probe(params, encodings, mask, n_heads: int, dropout: float = 0.0, use_LN=False, use_LN_input=False, use_softmax=True): +@bind(jax.jit, static_argnums=[4, 5, 6, 7, 8]) +def run_attention_probe( + dkey, params, encodings, mask, n_heads: int, dropout: float = 0.0, use_LN=False, use_LN_input=False, + use_softmax=True +): """ Runs full nonlinear attentive probe on input encodings (typically embedding vectors produced by some other model). Args: + dkey: JAX key for any internal noise to be applied + params: parameters tuple/list of probe encodings: input encoding vectors/data @@ -91,6 +98,8 @@ def run_attention_probe(params, encodings, mask, n_heads: int, dropout: float = use_LN: use layer normalization? + use_LN_input: use layer normalization on input encodings? + use_softmax: should softmax be applied to output of attention probe? (useful for classification) Returns: @@ -107,7 +116,7 @@ def run_attention_probe(params, encodings, mask, n_heads: int, dropout: float = if use_LN_input: learnable_query = layer_normalize(learnable_query, ln_in_mu, ln_in_scale) encodings = layer_normalize(encodings, ln_in_mu2, ln_in_scale2) - features = cross_attention(cross_attn_params, learnable_query, encodings, mask, n_heads, dropout) + features = cross_attention(dkey, cross_attn_params, learnable_query, encodings, mask, n_heads, dropout) # Perform a single self-attention block here # Self-Attention self_attn_params = (Wqs, bqs, Wks, bks, Wvs, bvs, Wouts, bouts) @@ -138,13 +147,15 @@ def run_attention_probe(params, encodings, mask, n_heads: int, dropout: float = outs = jax.nn.softmax(outs) return outs, features -@bind(jax.jit, static_argnums=[4, 5, 6, 7, 8]) -def eval_attention_probe(params, encodings, labels, mask, n_heads: int, dropout: float = 0.0, use_LN=False, use_LN_input=False, use_softmax=True): +@bind(jax.jit, static_argnums=[5, 6, 7, 8, 9]) +def eval_attention_probe(dkey, params, encodings, labels, mask, n_heads: int, dropout: float = 0.0, use_LN=False, use_LN_input=False, use_softmax=True): """ Runs and evaluates the nonlinear attentive probe given a paired set of encoding vectors and externally assigned labels/regression targets. Args: + dkey: JAX key to trigger any internal noise (as in drop-out) + params: parameters tuple/list of probe encodings: input encoding vectors/data @@ -165,7 +176,7 @@ def eval_attention_probe(params, encodings, labels, mask, n_heads: int, dropout: current loss value, output scores/probabilities """ # encodings: (B, hw, dim) - outs, _ = run_attention_probe(params, encodings, mask, n_heads, dropout, use_LN, use_LN_input, use_softmax) + outs, _ = run_attention_probe(dkey, params, encodings, mask, n_heads, dropout, use_LN, use_LN_input, use_softmax) if use_softmax: ## Multinoulli log likelihood for 1-of-K predictions L = -jnp.mean(jnp.sum(jnp.log(outs.clip(min=1e-5)) * labels, axis=1, keepdims=True)) else: ## MSE for real-valued outputs @@ -219,6 +230,7 @@ def __init__( self.use_softmax = use_softmax self.use_LN = use_LN self.use_LN_input = use_LN_input + self.dropout = 0.5 sigma = 0.05 ## cross-attention parameters @@ -275,42 +287,25 @@ def __init__( self.optim_params = adam.adam_init(self.probe_params) self.eta = 0.0002 #0.001 - def process(self, embedding_sequence): - """ - Runs the probe's inference scheme given an input batch of sequences of encodings/embeddings. - - Args: - embedding_sequence: a 3D tensor containing a batch of encoding sequences; shape (B, T, embed_dim) - - Returns: - probe output scores/probability values - """ - #print(embedding_sequence.shape) + def process(self, embeddings, dkey=None): + noise_key = None + if dkey is not None: + dkey, *subkeys = random.split(dkey, 2) + noise_key = subkeys[0] outs, feats = run_attention_probe( - self.probe_params, embedding_sequence, self.dev_mask, self.num_heads, 0.0, use_LN=self.use_LN, - use_LN_input=self.use_LN_input, use_softmax=self.use_softmax + noise_key, self.probe_params, embeddings, self.dev_mask, self.num_heads, 0.0, + use_LN=self.use_LN, use_LN_input=self.use_LN_input, use_softmax=self.use_softmax ) return outs - def update(self, embedding_sequence, labels, dkey=None): - """ - Runs and updates this probe given an input batch of sequences of encodings/embeddings and their externally - assigned labels/target vector values. - - Args: - embedding_sequence: a 3D tensor containing a batch of encoding sequences; shape (B, T, embed_dim) - - labels: target values that map to embedding sequence; shape (B, target_value_dim) - - Returns: - probe output scores/probability values - """ - # TODO: put in dkey to facilitate dropout - ## compute partial derivatives / adjustments to probe parameters - # NOTE: Viet: Change back to 0.0 for now for the code to run + def update(self, embeddings, labels, dkey=None): + noise_key = None + if dkey is not None: + dkey, *subkeys = random.split(dkey, 2) + noise_key = subkeys[0] outputs, grads = self.grad_fx( - self.probe_params, embedding_sequence, labels, self.mask, self.num_heads, dropout=0.0, use_LN=self.use_LN, - use_LN_input=self.use_LN_input, use_softmax=self.use_softmax + noise_key, self.probe_params, embeddings, labels, self.mask, self.num_heads, dropout=self.dropout, + use_LN=self.use_LN, use_LN_input=self.use_LN_input, use_softmax=self.use_softmax ) loss, predictions = outputs ## adjust parameters of probe diff --git a/ngclearn/utils/analysis/linear_probe.py b/ngclearn/utils/analysis/linear_probe.py index 370000f4..e6eb2a31 100644 --- a/ngclearn/utils/analysis/linear_probe.py +++ b/ngclearn/utils/analysis/linear_probe.py @@ -101,16 +101,7 @@ def __init__( self.optim_params = adam.adam_init(self.probe_params) self.eta = 0.001 - def process(self, embeddings): - """ - Runs the probe's inference scheme given an input batch of sequences of encodings/embeddings. - - Args: - embedding_sequence: a 3D tensor containing a batch of encoding sequences; shape (B, T, embed_dim) - - Returns: - probe output scores/probability values - """ + def process(self, embeddings, dkey=None): _embeddings = embeddings if len(_embeddings.shape) > 2: ## we flatten a sequence batch to 2D for a linear probe flat_dim = embeddings.shape[1] * embeddings.shape[2] @@ -118,19 +109,7 @@ def process(self, embeddings): outs = run_linear_probe(self.probe_params, _embeddings, use_softmax=self.use_softmax, use_LN=self.use_LN) return outs - def update(self, embeddings, labels): - """ - Runs and updates this probe given an input batch of sequences of encodings/embeddings and their externally - assigned labels/target vector values. - - Args: - embedding_sequence: a 3D tensor containing a batch of encoding sequences; shape (B, T, embed_dim) - - labels: target values that map to embedding sequence; shape (B, target_value_dim) - - Returns: - probe output scores/probability values - """ + def update(self, embeddings, labels, dkey=None): _embeddings = embeddings if len(_embeddings.shape) > 2: flat_dim = embeddings.shape[1] * embeddings.shape[2] diff --git a/ngclearn/utils/analysis/probe.py b/ngclearn/utils/analysis/probe.py index 7ad13238..531893d4 100644 --- a/ngclearn/utils/analysis/probe.py +++ b/ngclearn/utils/analysis/probe.py @@ -18,11 +18,36 @@ def __init__( self.batch_size = batch_size self.dev_batch_size = dev_batch_size - def process(self, embeddings): + def process(self, embeddings, dkey=None): + """ + Runs the probe's inference scheme given an input batch of sequences of encodings/embeddings. + + Args: + embeddings: a 3D tensor containing a batch of encoding sequences; shape (B, T, embed_dim) + + dkey: Optional JAX noise key + + Returns: + probe output scores/probability values + """ predictions = None return predictions - def update(self, embeddings, labels): + def update(self, embeddings, labels, dkey=None): + """ + Runs and updates this probe given an input batch of sequences of encodings/embeddings and their externally + assigned labels/target vector values. + + Args: + embeddings: a 3D tensor containing a batch of encoding sequences; shape (B, T, embed_dim) + + labels: target values that map to embedding sequence; shape (B, target_value_dim) + + dkey: Optional JAX noise key + + Returns: + probe output scores/probability values + """ L = predictions = None return L, predictions From 2feeced88f0a3f51d6e825628b77d65df8184295 Mon Sep 17 00:00:00 2001 From: Alexander Ororbia Date: Wed, 5 Mar 2025 20:15:55 -0500 Subject: [PATCH 48/61] cleaned up probes --- ngclearn/utils/analysis/probe.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/ngclearn/utils/analysis/probe.py b/ngclearn/utils/analysis/probe.py index 531893d4..0aa20946 100644 --- a/ngclearn/utils/analysis/probe.py +++ b/ngclearn/utils/analysis/probe.py @@ -79,7 +79,7 @@ def predict(self, data, batch_size=None): x_mb = _data[s_ptr:e_ptr, :, :] ## slice out 3D batch tensor s_ptr = e_ptr e_ptr += x_mb.shape[0] - y_mu = self.process(x_mb) + y_mu = self.process(x_mb, dkey=None) Y_mu.append(y_mu) Y_mu = jnp.concatenate(Y_mu, axis=0) return Y_mu @@ -143,8 +143,9 @@ def fit(self, dataset, dev_dataset=None, n_iter=50, patience=20): s_ptr = e_ptr e_ptr += x_mb.shape[0] Ns += x_mb.shape[0] + self.dkey, *subkeys = random.split(self.dkey, 2) - _L, py = self.update(x_mb, y_mb) + _L, py = self.update(x_mb, y_mb, dkey=subkeys[0]) acc = jnp.sum(jnp.equal(jnp.argmax(py, axis=1), jnp.argmax(y_mb, axis=1))) + acc L = (_L * x_mb.shape[0]) + L ## we remove the batch division from loss w.r.t. x_mb/y_mb From 56f006cc4d19f4b000e6a18b824aa10670ee850f Mon Sep 17 00:00:00 2001 From: Alexander Ororbia Date: Wed, 5 Mar 2025 20:18:32 -0500 Subject: [PATCH 49/61] cleaned up probes --- ngclearn/utils/analysis/attentive_probe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ngclearn/utils/analysis/attentive_probe.py b/ngclearn/utils/analysis/attentive_probe.py index 7e17e953..8c11c613 100644 --- a/ngclearn/utils/analysis/attentive_probe.py +++ b/ngclearn/utils/analysis/attentive_probe.py @@ -282,7 +282,7 @@ def __init__( self.probe_params = (learnable_query, *cross_attn_params, *self_attn_params, *mlp_params, *ln_in_params) ## set up gradient calculator - self.grad_fx = jax.value_and_grad(eval_attention_probe, argnums=0, has_aux=True) + self.grad_fx = jax.value_and_grad(eval_attention_probe, argnums=0, has_aux=True, allow_int=True) ## set up update rule/optimizer self.optim_params = adam.adam_init(self.probe_params) self.eta = 0.0002 #0.001 From 1b7bff8fe175466c7bfd074b8ab7579876bf4b99 Mon Sep 17 00:00:00 2001 From: Alexander Ororbia Date: Wed, 5 Mar 2025 20:19:24 -0500 Subject: [PATCH 50/61] cleaned up probes --- ngclearn/utils/analysis/attentive_probe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ngclearn/utils/analysis/attentive_probe.py b/ngclearn/utils/analysis/attentive_probe.py index 8c11c613..d8301c2e 100644 --- a/ngclearn/utils/analysis/attentive_probe.py +++ b/ngclearn/utils/analysis/attentive_probe.py @@ -282,7 +282,7 @@ def __init__( self.probe_params = (learnable_query, *cross_attn_params, *self_attn_params, *mlp_params, *ln_in_params) ## set up gradient calculator - self.grad_fx = jax.value_and_grad(eval_attention_probe, argnums=0, has_aux=True, allow_int=True) + self.grad_fx = jax.value_and_grad(eval_attention_probe, argnums=1, has_aux=True) #, allow_int=True) ## set up update rule/optimizer self.optim_params = adam.adam_init(self.probe_params) self.eta = 0.0002 #0.001 From f38373f44b3033818373cbdf6564263fa2a8308d Mon Sep 17 00:00:00 2001 From: Alexander Ororbia Date: Wed, 5 Mar 2025 20:21:24 -0500 Subject: [PATCH 51/61] generalized dropout in terms of shape --- ngclearn/utils/model_utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/ngclearn/utils/model_utils.py b/ngclearn/utils/model_utils.py index e5d8fcd5..2c245dfe 100755 --- a/ngclearn/utils/model_utils.py +++ b/ngclearn/utils/model_utils.py @@ -567,8 +567,7 @@ def drop_out(dkey, input, rate=0.0): Returns: output as well as binary mask """ - eps = random.uniform(dkey, (input.shape[0],input.shape[1]), - minval=0.0, maxval=1.0) + eps = random.uniform(dkey, shape=input.shape, minval=0.0, maxval=1.0) mask = (eps <= (1.0 - rate)).astype(jnp.float32) mask = mask * (1.0 / (1.0 - rate)) ## apply inverted dropout scheme output = input * mask From 012395b984de22a55d79aff5186a0f555dbac74c Mon Sep 17 00:00:00 2001 From: Alexander Ororbia Date: Wed, 5 Mar 2025 20:24:29 -0500 Subject: [PATCH 52/61] tweak to atten probe --- ngclearn/utils/analysis/attentive_probe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ngclearn/utils/analysis/attentive_probe.py b/ngclearn/utils/analysis/attentive_probe.py index d8301c2e..4812aa03 100644 --- a/ngclearn/utils/analysis/attentive_probe.py +++ b/ngclearn/utils/analysis/attentive_probe.py @@ -70,7 +70,7 @@ def cross_attention(dkey, params: tuple, x1: jax.Array, x2: jax.Array, mask: jax score = jax.nn.softmax(score, axis=-1) # (B, H, T, S) score = score.astype(q.dtype) # (B, H, T, S) if dropout_rate > 0.: - score = drop_out(dkey, input=score, rate=dropout_rate) ## NOTE: normally you apply dropout here + score, _ = drop_out(dkey, input=score, rate=dropout_rate) ## NOTE: normally you apply dropout here attention = jnp.einsum("BHTS,BHSE->BHTE", score, v) # (B, T, H, E) attention = attention.transpose([0, 2, 1, 3]).reshape((B, T, -1)) # (B, T, H, E) => (B, T, D) return attention @ Wout + bout # (B, T, Dq) From 53ed77314fd6e0e9ef4bad21eb9ad30cc943af32 Mon Sep 17 00:00:00 2001 From: Alexander Ororbia Date: Wed, 5 Mar 2025 20:27:45 -0500 Subject: [PATCH 53/61] tweak to atten probe --- ngclearn/utils/analysis/attentive_probe.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ngclearn/utils/analysis/attentive_probe.py b/ngclearn/utils/analysis/attentive_probe.py index 4812aa03..36bde6fd 100644 --- a/ngclearn/utils/analysis/attentive_probe.py +++ b/ngclearn/utils/analysis/attentive_probe.py @@ -123,7 +123,7 @@ def run_attention_probe( skip = features if use_LN: features = layer_normalize(features, Wlnattn_mu, Wlnattn_scale) - features = cross_attention(self_attn_params, features, features, None, n_heads, dropout) + features = cross_attention(dkey, self_attn_params, features, features, None, n_heads, dropout) features = features + skip features = features[:, 0] # (B, 1, dim) => (B, dim) # MLP From 1fbbf93334f64906cb307d87e30853159c0ac5e7 Mon Sep 17 00:00:00 2001 From: Alexander Ororbia Date: Thu, 6 Mar 2025 14:12:29 -0500 Subject: [PATCH 54/61] added silu/swish/elu to model_utils --- ngclearn/utils/model_utils.py | 76 +++++++++++++++++++++++++++++++---- 1 file changed, 68 insertions(+), 8 deletions(-) diff --git a/ngclearn/utils/model_utils.py b/ngclearn/utils/model_utils.py index 2c245dfe..ffb624d7 100755 --- a/ngclearn/utils/model_utils.py +++ b/ngclearn/utils/model_utils.py @@ -453,6 +453,47 @@ def inverse_logistic(x, clip_bound=0.03): # 0.03 x_ = jnp.clip(x_, clip_bound, 1.0 - clip_bound) return jnp.log( x_/((1.0 - x_) + 1e-6) ) +@jit +def swish(x, beta): + """ + Applies the Swish parameterized activation, proposed in Ramachandran et al., 2017 + ("Searching for Activation Functions"). + + Args: + x: data to transform via inverse logistic function + + beta: coefficient/parameters to weight input x by + + Returns: + output of the Swish activation + """ + return x * sigmoid(x * beta) + +@jit +def d_swish(x, beta): + # df/dx = beta * [ 1/(exp(-x) + 1) + (exp(-x) * x) / (exp(-x) + 1)^2] + # df/dx = beta * sigmoid(x * beta) * (1 - sigmoid(x) * beta) + exp_neg_x = jnp.exp(-x) + _x = (1./(exp_neg_x + 1.)) + (exp_neg_x * x)/jnp.square(exp_neg_x+1) + return _x * beta + +@jit +def silu(x): + """ + Applies the sigmoid-weighted linear unit (SiLU or SiL) activation. + + Args: + x: data to transform via inverse logistic function + + Returns: + output of the Swish activation + """ + return swish(x, beta=1.) + +@jit +def d_silu(x): + return d_swish(x, beta=1.) + @jit def gelu(x): """ @@ -464,14 +505,33 @@ def gelu(x): Returns: output of the GeLU activation """ - return x * sigmoid(x * 1.702) ## approximate GeLU + return swish(x, beta=1.702) ## approximate GeLU # beta=1.4 @jit def d_gelu(x): # df/dx = 1.702 * [ 1/(exp(-x) + 1) + (exp(-x) * x) / (exp(-x) + 1)^2] - exp_neg_x = jnp.exp(-x) - _x = (1./(exp_neg_x + 1.)) + (exp_neg_x * x)/jnp.square(exp_neg_x+1) - return _x * 1.702 + return d_swish(x, beta=1.702) # beta=1.4 + +@jit +def elu(x, alpha=1.): + """ + Applies the exponential linear unit (ELU) activation. + + Args: + x: data to transform via inverse logistic function + + alpha: coefficient/parameters to weight input x by + + Returns: + output of the GeLU activation + """ + mask = x >= 0. + return x * mask + ((jnp.exp(x) - 1) * alpha) * (1. - mask) + +@jit +def elu(x, alpha=1.): + mask = (x >= 0.) + return mask + (1. - mask) * (jnp.exp(x) * alpha) @jit def softmax(x, tau=0.0): @@ -553,24 +613,24 @@ def layer_normalize(x, shift=0., scale=1.): return _x * scale + shift @jit -def drop_out(dkey, input, rate=0.0): +def drop_out(dkey, data, rate=0.0): """ Applies a drop-out transform to an input matrix. Args: dkey: Jax randomness key for this operator - input: data to apply random/drop-out mask to + data: input data to apply random/drop-out mask to rate: probability of a dimension being dropped Returns: output as well as binary mask """ - eps = random.uniform(dkey, shape=input.shape, minval=0.0, maxval=1.0) + eps = random.uniform(dkey, shape=data.shape, minval=0.0, maxval=1.0) mask = (eps <= (1.0 - rate)).astype(jnp.float32) mask = mask * (1.0 / (1.0 - rate)) ## apply inverted dropout scheme - output = input * mask + output = data * mask return output, mask From 23e8c844d846cd2d925eadbd904ef9725f4e6215 Mon Sep 17 00:00:00 2001 From: Alexander Ororbia Date: Thu, 6 Mar 2025 14:26:39 -0500 Subject: [PATCH 55/61] cleaned up model_utils --- ngclearn/utils/model_utils.py | 43 ++++++++++++++++++++++------------- 1 file changed, 27 insertions(+), 16 deletions(-) diff --git a/ngclearn/utils/model_utils.py b/ngclearn/utils/model_utils.py index ffb624d7..d3e003d8 100755 --- a/ngclearn/utils/model_utils.py +++ b/ngclearn/utils/model_utils.py @@ -72,18 +72,20 @@ def create_function(fun_name, args=None): Activation function creation routine. Args: - fun_name: string name of activation function to produce - (Currently supports: "tanh", "relu", "lrelu", "identity") + fun_name: string name of activation function to produce; + Currently supports: "tanh", "bkwta" (binary K-winners-take-all), "sigmoid", "relu", "lrelu", "relu6", + "elu", "silu", "gelu", "softplus", "softmax" (derivative not supported), "unit_threshold", "heaviside", + "identity" Returns: function fx, first derivative of function (w.r.t. input) dfx """ - fx = None - dfx = None + fx = None ## the function + dfx = None ## the first derivative of function w.r.t. its input if fun_name == "tanh": fx = tanh dfx = d_tanh - elif "kwta" in fun_name: + elif "bkwta" in fun_name: fx = bkwta dfx = bkwta #d_identity elif fun_name == "sigmoid": @@ -98,6 +100,15 @@ def create_function(fun_name, args=None): elif fun_name == "relu6": fx = relu6 dfx = d_relu6 + elif fun_name == "elu": + fx = elu + dfx = d_elu + elif fun_name == "silu": + fx = silu + dfx = d_silu + elif fun_name == "gelu": + fx = gelu + dfx = d_gelu elif fun_name == "softplus": fx = softplus dfx = d_softplus @@ -127,35 +138,35 @@ def bkwta(x, nWTA=5): #5 10 15 #K=50): return topK @partial(jit, static_argnums=[2, 3, 4]) -def normalize_matrix(M, wnorm, order=1, axis=0, scale=1.): +def normalize_matrix(data, wnorm, order=1, axis=0, scale=1.): """ Normalizes the values in matrix to have a particular norm across each vector span. Args: - M: (2D) matrix to normalize + data: (2D) data matrix to normalize - wnorm: target norm for each + wnorm: target norm for each row/column of data matrix order: order of norm to use in normalization (Default: 1); note that `ord=1` results in the L1-norm, `ord=2` results in the L2-norm axis: 0 (apply to column vectors), 1 (apply to row vectors) - scale: step modifier to produce the projected matrix + scale: step modifier to produce the projected matrix (Unused) Returns: a normalized value matrix """ if order == 2: ## denominator is L2 norm - wOrdSum = jnp.maximum(jnp.sqrt(jnp.sum(jnp.square(M), axis=axis, keepdims=True)), 1e-8) + wOrdSum = jnp.maximum(jnp.sqrt(jnp.sum(jnp.square(data), axis=axis, keepdims=True)), 1e-8) else: ## denominator is L1 norm - wOrdSum = jnp.maximum(jnp.sum(jnp.abs(M), axis=axis, keepdims=True), 1e-8) + wOrdSum = jnp.maximum(jnp.sum(jnp.abs(data), axis=axis, keepdims=True), 1e-8) m = (wOrdSum == 0.).astype(dtype=jnp.float32) wOrdSum = wOrdSum * (1. - m) + m #wAbsSum[wAbsSum == 0.] = 1. - _M = M * (wnorm/wOrdSum) - #dM = ((wnorm/wOrdSum) - 1.) * M - #_M = M + dM * scale - return _M + _data = data * (wnorm/wOrdSum) + #d_data = ((wnorm/wOrdSum) - 1.) * data + #_data = data + d_data * scale + return _data @jit def clamp_min(x, min_val): @@ -529,7 +540,7 @@ def elu(x, alpha=1.): return x * mask + ((jnp.exp(x) - 1) * alpha) * (1. - mask) @jit -def elu(x, alpha=1.): +def d_elu(x, alpha=1.): mask = (x >= 0.) return mask + (1. - mask) * (jnp.exp(x) * alpha) From 695e9d8257257daec69a47b0cfc9935641eb814c Mon Sep 17 00:00:00 2001 From: Viet Dung Nguyen Date: Fri, 7 Mar 2025 00:26:35 -0500 Subject: [PATCH 56/61] fix bug in attention probe dropout, fix bug in None noise_key passed in the probing jit function, add the spliting of noise_keys to two dropout in two cross attention --- ngclearn/utils/analysis/attentive_probe.py | 19 +++++++++++++------ 1 file changed, 13 insertions(+), 6 deletions(-) diff --git a/ngclearn/utils/analysis/attentive_probe.py b/ngclearn/utils/analysis/attentive_probe.py index 36bde6fd..0129a267 100644 --- a/ngclearn/utils/analysis/attentive_probe.py +++ b/ngclearn/utils/analysis/attentive_probe.py @@ -70,7 +70,7 @@ def cross_attention(dkey, params: tuple, x1: jax.Array, x2: jax.Array, mask: jax score = jax.nn.softmax(score, axis=-1) # (B, H, T, S) score = score.astype(q.dtype) # (B, H, T, S) if dropout_rate > 0.: - score, _ = drop_out(dkey, input=score, rate=dropout_rate) ## NOTE: normally you apply dropout here + score, _ = drop_out(dkey, score, rate=dropout_rate) ## NOTE: normally you apply dropout here attention = jnp.einsum("BHTS,BHSE->BHTE", score, v) # (B, T, H, E) attention = attention.transpose([0, 2, 1, 3]).reshape((B, T, -1)) # (B, T, H, E) => (B, T, D) return attention @ Wout + bout # (B, T, Dq) @@ -105,6 +105,8 @@ def run_attention_probe( Returns: output scores/probabilities, cross-attention (hidden) features """ + # Two separate dkeys for each dropout in two cross attention + dkey1, dkey2 = random.split(dkey, 2) # encoded_image_feature: (B, hw, dim) #learnable_query, *_params) = params learnable_query, Wq, bq, Wk, bk, Wv, bv, Wout, bout,\ @@ -116,14 +118,14 @@ def run_attention_probe( if use_LN_input: learnable_query = layer_normalize(learnable_query, ln_in_mu, ln_in_scale) encodings = layer_normalize(encodings, ln_in_mu2, ln_in_scale2) - features = cross_attention(dkey, cross_attn_params, learnable_query, encodings, mask, n_heads, dropout) + features = cross_attention(dkey1, cross_attn_params, learnable_query, encodings, mask, n_heads, dropout) # Perform a single self-attention block here # Self-Attention self_attn_params = (Wqs, bqs, Wks, bks, Wvs, bvs, Wouts, bouts) skip = features if use_LN: features = layer_normalize(features, Wlnattn_mu, Wlnattn_scale) - features = cross_attention(dkey, self_attn_params, features, features, None, n_heads, dropout) + features = cross_attention(dkey2, self_attn_params, features, features, None, n_heads, dropout) features = features + skip features = features[:, 0] # (B, 1, dim) => (B, dim) # MLP @@ -222,7 +224,7 @@ def __init__( super().__init__(dkey, batch_size, **kwargs) assert attn_dim % num_heads == 0, f"`attn_dim` must be divisible by `num_heads`. Got {attn_dim} and {num_heads}." assert learnable_query_dim % num_heads == 0, f"`learnable_query_dim` must be divisible by `num_heads`. Got {learnable_query_dim} and {num_heads}." - self.dkey, *subkeys = random.split(self.dkey, 25) + self.dkey, *subkeys = random.split(self.dkey, 26) self.num_heads = num_heads self.source_seq_length = source_seq_length self.input_dim = input_dim @@ -287,8 +289,12 @@ def __init__( self.optim_params = adam.adam_init(self.probe_params) self.eta = 0.0002 #0.001 + # Finally, the dkey for the noise_key + self.noise_key = subkeys[24] + def process(self, embeddings, dkey=None): - noise_key = None + # noise_key = None + noise_key = self.noise_key if dkey is not None: dkey, *subkeys = random.split(dkey, 2) noise_key = subkeys[0] @@ -299,7 +305,8 @@ def process(self, embeddings, dkey=None): return outs def update(self, embeddings, labels, dkey=None): - noise_key = None + # noise_key = None + noise_key = self.noise_key if dkey is not None: dkey, *subkeys = random.split(dkey, 2) noise_key = subkeys[0] From 04e1343096bbda6e0ab96e609dcf6e070effa5be Mon Sep 17 00:00:00 2001 From: Viet Nguyen Date: Sun, 9 Mar 2025 20:48:00 -0400 Subject: [PATCH 57/61] hyperparameter tunning arguments added --- ngclearn/utils/analysis/attentive_probe.py | 15 +++++++++++---- ngclearn/utils/analysis/probe.py | 9 +++++++-- 2 files changed, 18 insertions(+), 6 deletions(-) diff --git a/ngclearn/utils/analysis/attentive_probe.py b/ngclearn/utils/analysis/attentive_probe.py index 0129a267..f721d69a 100644 --- a/ngclearn/utils/analysis/attentive_probe.py +++ b/ngclearn/utils/analysis/attentive_probe.py @@ -219,7 +219,9 @@ class AttentiveProbe(Probe): """ def __init__( self, dkey, source_seq_length, input_dim, out_dim, num_heads=8, attn_dim=64, - target_seq_length=1, learnable_query_dim=32, batch_size=1, hid_dim=32, use_LN=True, use_LN_input=False, use_softmax=True, **kwargs + target_seq_length=1, learnable_query_dim=32, batch_size=1, hid_dim=32, + use_LN=True, use_LN_input=False, use_softmax=True, dropout=0.5, eta=0.0002, + eta_decay=0.0, min_eta=1e-5, **kwargs ): super().__init__(dkey, batch_size, **kwargs) assert attn_dim % num_heads == 0, f"`attn_dim` must be divisible by `num_heads`. Got {attn_dim} and {num_heads}." @@ -232,9 +234,9 @@ def __init__( self.use_softmax = use_softmax self.use_LN = use_LN self.use_LN_input = use_LN_input - self.dropout = 0.5 + self.dropout = dropout - sigma = 0.05 + sigma = 0.02 ## cross-attention parameters Wq = random.normal(subkeys[0], (learnable_query_dim, attn_dim)) * sigma bq = random.normal(subkeys[1], (1, attn_dim)) * sigma @@ -287,7 +289,10 @@ def __init__( self.grad_fx = jax.value_and_grad(eval_attention_probe, argnums=1, has_aux=True) #, allow_int=True) ## set up update rule/optimizer self.optim_params = adam.adam_init(self.probe_params) - self.eta = 0.0002 #0.001 + # Learning rate scheduling + self.eta = eta #0.001 + self.eta_decay = eta_decay + self.min_eta = min_eta # Finally, the dkey for the noise_key self.noise_key = subkeys[24] @@ -319,5 +324,7 @@ def update(self, embeddings, labels, dkey=None): self.optim_params, self.probe_params = adam.adam_step( self.optim_params, self.probe_params, grads, eta=self.eta ) + + self.eta = max(self.min_eta, self.eta - self.eta_decay * self.eta) return loss, predictions diff --git a/ngclearn/utils/analysis/probe.py b/ngclearn/utils/analysis/probe.py index 0aa20946..165ab5f4 100644 --- a/ngclearn/utils/analysis/probe.py +++ b/ngclearn/utils/analysis/probe.py @@ -150,9 +150,14 @@ def fit(self, dataset, dev_dataset=None, n_iter=50, patience=20): L = (_L * x_mb.shape[0]) + L ## we remove the batch division from loss w.r.t. x_mb/y_mb if dev_data is not None: - print(f"\r{ii} L = {L / Ns:.3f} Acc = {acc / Ns:.2f} Dev.Acc = {best_acc:.2f}", end="") + print_string = f"\r{ii} L = {L / Ns:.3f} Acc = {acc / Ns:.2f} Dev.Acc = {best_acc:.2f}" else: - print(f"\r{ii} L = {L / Ns:.3f} Acc = {acc / Ns:.2f}", end="") + print_string = f"\r{ii} L = {L / Ns:.3f} Acc = {acc / Ns:.2f}" + + if hasattr(self, "eta"): + print_string += f" LR = {getattr(self, 'eta'):.6f}" + + print(print_string, end = "") acc = acc / Ns L = L / Ns ## compute current loss over (train) dataset From 7bfd8acd47704deabcd9d6de628a27492431b5c4 Mon Sep 17 00:00:00 2001 From: Viet Dung Nguyen Date: Wed, 12 Mar 2025 19:07:13 -0400 Subject: [PATCH 58/61] remove unused local variables --- ngclearn/utils/model_utils.py | 1 - 1 file changed, 1 deletion(-) diff --git a/ngclearn/utils/model_utils.py b/ngclearn/utils/model_utils.py index fd3a9eda..a55dc580 100755 --- a/ngclearn/utils/model_utils.py +++ b/ngclearn/utils/model_utils.py @@ -91,7 +91,6 @@ def create_function(fun_name, args=None): elif fun_name == "sine": fx = sine dfx = d_sine - omega_0 = args elif fun_name == "sigmoid": fx = sigmoid dfx = d_sigmoid From 27ae7e20042bb684b97270b6f5359fa2c2c22555 Mon Sep 17 00:00:00 2001 From: Viet Dung Nguyen Date: Wed, 12 Mar 2025 19:20:44 -0400 Subject: [PATCH 59/61] update note --- ngclearn/utils/model_utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/ngclearn/utils/model_utils.py b/ngclearn/utils/model_utils.py index a55dc580..68a1cb9e 100755 --- a/ngclearn/utils/model_utils.py +++ b/ngclearn/utils/model_utils.py @@ -590,11 +590,12 @@ def softmax(x, tau=0.0): Returns: a (N x D) probability distribution output block """ + # TODO: Do we need to also consider for edge case the division of tau by 0 is invalid? if tau > 0.0: x = x / tau max_x = jnp.max(x, axis=1, keepdims=True) exp_x = jnp.exp(x - max_x) - return exp_x / jnp.sum(exp_x, axis=1, keepdims=True) + return exp_x / jnp.sum(exp_x, axis=1, keepdims=True) # TODO: We also have to take care of the case where the sum is 0 @jit def threshold_soft(x, lmbda): From 92633f9768a14778bd43a3f597d4063847b4d740 Mon Sep 17 00:00:00 2001 From: Viet Dung Nguyen Date: Wed, 12 Mar 2025 20:04:53 -0400 Subject: [PATCH 60/61] update model utils --- ngclearn/utils/model_utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ngclearn/utils/model_utils.py b/ngclearn/utils/model_utils.py index 68a1cb9e..028de1de 100755 --- a/ngclearn/utils/model_utils.py +++ b/ngclearn/utils/model_utils.py @@ -85,7 +85,7 @@ def create_function(fun_name, args=None): if fun_name == "tanh": fx = tanh dfx = d_tanh - elif "bkwta" in fun_name: + elif fun_name == "bkwta": fx = bkwta dfx = bkwta #d_identity elif fun_name == "sine": @@ -651,8 +651,8 @@ def layer_normalize(x, shift=0., scale=1.): layer-normalized data samples `x` """ xmu = jnp.mean(x, axis=1, keepdims=True) - xsigma = jnp.sqrt(jnp.mean(jnp.square(x - xmu)).clip(min=1e-6)) - _x = (x - xmu)/(xsigma + 1e-6) + xsigma = jnp.sqrt(jnp.mean(jnp.square(x - xmu)).clip(min=1e-6)).clip(min=1e-6) + _x = (x - xmu) / xsigma return _x * scale + shift @jit From 08b4d12225c85e81036703764e6dc4893daec0ac Mon Sep 17 00:00:00 2001 From: Viet Dung Nguyen Date: Wed, 12 Mar 2025 20:08:22 -0400 Subject: [PATCH 61/61] remove notes --- ngclearn/utils/model_utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/ngclearn/utils/model_utils.py b/ngclearn/utils/model_utils.py index 028de1de..7308a1d0 100755 --- a/ngclearn/utils/model_utils.py +++ b/ngclearn/utils/model_utils.py @@ -590,12 +590,11 @@ def softmax(x, tau=0.0): Returns: a (N x D) probability distribution output block """ - # TODO: Do we need to also consider for edge case the division of tau by 0 is invalid? if tau > 0.0: x = x / tau max_x = jnp.max(x, axis=1, keepdims=True) exp_x = jnp.exp(x - max_x) - return exp_x / jnp.sum(exp_x, axis=1, keepdims=True) # TODO: We also have to take care of the case where the sum is 0 + return exp_x / jnp.sum(exp_x, axis=1, keepdims=True) @jit def threshold_soft(x, lmbda):