rate_loss

Contents

rate_loss#

class braintools.fit.losses.rate_loss(prediction, target, dt=None, *, per_neuron=False)#

Compute the squared difference of mean firing rates (in Hz).

Each input is a dense spike raster of shape (T,) or (T, N) (time on axis 0) or a SpikeTrains. The rate of a neuron is its spike count divided by the recording length: T * dt for a raster, t_stop - t_start for a SpikeTrains. The loss is

\[L = (\bar r_{\mathrm{pred}} - \bar r_{\mathrm{target}})^2,\]

the squared difference of the population-mean rates, or with per_neuron=True

\[L = \frac{1}{N} \sum_{n=1}^{N} (r_{\mathrm{pred},n} - r_{\mathrm{target},n})^2.\]
Parameters:

  • prediction (Any) – Model spikes. Raster values may be real (e.g. surrogate-gradient spikes); the count is their sum.

  • target (Any) – Data spikes. Rasters may have different T from the prediction.

  • dt (Any) – Raster time step. A plain number is in milliseconds (the spike-time convention of braintools.analysis.spike and brainstate), not seconds as in psd_log_distance(); pass a Quantity (e.g. 0.1 * u.ms) to avoid ambiguity. Defaults to brainstate.environ.get_dt(). Not used for SpikeTrains inputs.

  • per_neuron (bool) – If False, compare the population-mean rates. If True, average the squared per-neuron rate differences (both inputs need the same N).

Returns:

Scalar loss in Hz².

Return type:

Array

Raises:

  • ValueError – If a raster is not 1-D or 2-D, dt is missing for a raster, or per_neuron=True with different neuron counts.

  • TypeError – If a raster carries units or dt is not a time.

See also

cv_loss, van_rossum_loss, braintools.analysis.spike.firing_rate

Notes

Works under jax.jit and jax.vmap (dt may be traced) and is differentiable with respect to raster values, so it trains spiking networks through surrogate spikes. A SpikeTrains rate does not depend on the spike times, so it carries no gradient.

Examples

>>> import brainunit as u
>>> import jax.numpy as jnp
>>> from braintools.fit import losses
>>> from braintools.spike import SpikeTrains
>>> pred = jnp.zeros((1000, 4)).at[::100].set(1.)   # 10 Hz at dt = 1 ms
>>> target = jnp.zeros((1000, 4)).at[::50].set(1.)  # 20 Hz
>>> float(losses.rate_loss(pred, target, dt=1. * u.ms))
100.0
>>> st = SpikeTrains.from_dense(target, dt=1. * u.ms)
>>> float(losses.rate_loss(pred, st, dt=1.))
100.0
Contents