snntorch.weight_init

Spiking-aware weight initialization for snnTorch.

Provides fluctuation-driven initialization methods that account for neuron membrane time constants, synaptic time constants, and presynaptic firing rates.

Based on:: Rossbroich, J., Gygax, J. & Zenke, F. (2022). Fluctuation-driven initialization for spiking neural network training. Neuromorphic Computing and Engineering, 2(4), 044016. https://iopscience.iop.org/article/10.1088/2634-4386/ac97bb

snntorch.weight_init.compute_psp_kernel_integrals(tau_mem: float, tau_syn: float, dt: float = 0.001, mode: str = 'analytical') → Tuple[float, float][source]

Compute the integrals of the PSP kernel (epsilon_bar and epsilon_hat).

These integrals are used by the fluctuation-driven initialization to compute weight parameters from target membrane statistics.

Args:

tau_mem: Membrane time constant in seconds. tau_syn: Synaptic time constant in seconds. dt: Simulation timestep in seconds. Only used when mode=”numerical”. mode: “analytical” or “numerical”. Default: “analytical”.

Returns:

Tuple of (epsilon_bar, epsilon_hat).

Raises:

ValueError: If tau_mem or tau_syn are not positive. ValueError: If mode is not “analytical” or “numerical”.

Example:

>>> from snntorch import weight_init
>>> ebar, ehat = weight_init.compute_psp_kernel_integrals(
...     tau_mem=0.02, tau_syn=0.005
... )
>>> print(f"epsilon_bar={ebar:.4f}, epsilon_hat={ehat:.6f}")

snntorch.weight_init.fluctuation_driven_normal_(tensor: Tensor, tau_mem: float, tau_syn: float, nu: float, mu_u: float = 0.0, sigma_u: float = 1.0, theta: float = 1.0, dt: float = 0.001, mode: str = 'analytical', generator: Generator | None = None) → Tensor[source]

Initialize weights from a normal distribution using fluctuation-driven initialization.

This method targets a membrane potential with mean mu_u and standard deviation sigma_u, placing neurons near threshold for optimal surrogate gradient learning.

Args:

tensor: Weight tensor to initialize (modified in-place). tau_mem: Membrane time constant in seconds. tau_syn: Synaptic time constant in seconds. nu: Estimated presynaptic firing rate in Hz. mu_u: Target mean membrane potential. Default: 0.0 (balanced). sigma_u: Target standard deviation of membrane potential. Default: 1.0. theta: Spike threshold. Default: 1.0. dt: Simulation timestep in seconds. Default: 1e-3. mode: “analytical” or “numerical” for PSP kernel integrals.

Default: “analytical”.

generator: Optional torch.Generator for reproducibility.

Returns:

The initialized tensor (same object, modified in-place).

Raises:

ValueError: If computed sigma_w is negative. ValueError: If tensor has fewer than 2 dimensions. ValueError: If mode is not “analytical” or “numerical”. ValueError: If any parameter is not positive.

Example:

>>> import torch.nn as nn
>>> import snntorch as snn
>>> fc = nn.Linear(784, 128)
>>> snn.weight_init.fluctuation_driven_normal_(
...     fc.weight, tau_mem=0.02, tau_syn=0.005, nu=10.0, sigma_u=1.0
... )

snntorch.weight_init.fluctuation_driven_normal_centered_(tensor: Tensor, tau_mem: float, tau_syn: float, nu: float, sigma_u: float = 1.0, dt: float = 0.001, mode: str = 'analytical', generator: Generator | None = None) → Tensor[source]

Initialize weights using centered fluctuation-driven initialization (mu_u=0).

This is a simplified version assuming balanced excitation/inhibition where the target mean membrane potential is zero. This is the most commonly used variant.

The weight distribution is N(0, sigma_w) where:: sigma_w = sigma_u / sqrt(n * nu * epsilon_hat)

(Eq. 9 from Rossbroich et al. 2022)

Args:

tensor: Weight tensor to initialize (modified in-place). tau_mem: Membrane time constant in seconds. tau_syn: Synaptic time constant in seconds. nu: Estimated presynaptic firing rate in Hz. sigma_u: Target standard deviation of membrane potential. Default: 1.0. dt: Simulation timestep in seconds. Default: 1e-3. mode: “analytical” or “numerical” for PSP kernel integrals.

Default: “analytical”.

generator: Optional torch.Generator for reproducibility.

Returns:

The initialized tensor (same object, modified in-place).

Raises:

ValueError: If tensor has fewer than 2 dimensions. ValueError: If mode is not “analytical” or “numerical”. ValueError: If any parameter is not positive.

Example:

>>> import torch.nn as nn
>>> import snntorch as snn
>>> conv = nn.Conv2d(64, 128, 3, padding=1)
>>> snn.weight_init.fluctuation_driven_normal_centered_(
...     conv.weight, tau_mem=0.02, tau_syn=0.005, nu=15.0
... )

Overview

Weight initialization functions tailored for spiking neural networks. These functions initialize synaptic weights based on the neuron’s membrane time constant (tau_mem), synaptic time constant (tau_syn), and presynaptic firing rate (nu), placing neurons in the fluctuation-driven optimal regime for surrogate gradient learning.

The fluctuation-driven regime places neurons near (but below) the spike threshold, where they are most sensitive to input changes. This initialization method accounts for SNN-specific neuron dynamics, unlike standard initializers (e.g., Kaiming, Xavier) which assume rate-based activations.

Based on:: Rossbroich, J., Gygax, J. & Zenke, F. (2022). Fluctuation-driven initialization for spiking neural network training. Neuromorphic Computing and Engineering, 2(4), 044016.

Usage Example

import torch
import torch.nn as nn
import snntorch as snn
from snntorch import weight_init

# Create a simple SNN
net = nn.Sequential(
    nn.Linear(784, 256),
    snn.Leaky(beta=0.95),
    nn.Linear(256, 128),
    snn.Leaky(beta=0.95),
    nn.Linear(128, 10),
    snn.Leaky(beta=0.95),
)

# Apply fluctuation-driven initialization to all Linear layers
for layer in net:
    if isinstance(layer, nn.Linear):
        weight_init.fluctuation_driven_normal_centered_(
            layer.weight,
            tau_mem=0.02,   # 20ms membrane time constant
            tau_syn=0.005,  # 5ms synaptic time constant
            nu=10.0,        # 10 Hz estimated firing rate
            sigma_u=1.0,    # Target membrane std
        )

Mathematical Background

The membrane potential of a current-based LIF neuron driven by Poisson inputs follows a Gaussian distribution with mean mu_u and variance sigma_u^2:

\[ \begin{align}\begin{aligned}\mu_U = n \cdot \mu_W \cdot \nu \cdot \bar{\epsilon}\\\sigma_U^2 = n \cdot (\sigma_W^2 + \mu_W^2) \cdot \nu \cdot \hat{\epsilon}\end{aligned}\end{align} \]

where:

n is the number of presynaptic inputs (fan-in)
nu is the presynaptic firing rate (Hz)
epsilon_bar is the integral of the PSP kernel
epsilon_hat is the integral of the squared PSP kernel

The initialization inverts these equations to compute weight parameters:

\[ \begin{align}\begin{aligned}\mu_W = \frac{\mu_U}{n \cdot \nu \cdot \bar{\epsilon}}\\\sigma_W = \sqrt{\frac{\sigma_U^2}{n \cdot \nu \cdot \hat{\epsilon}} - \mu_W^2}\end{aligned}\end{align} \]

For the centered case (balanced excitation/inhibition, mu_u = 0):

\[\sigma_W = \frac{\sigma_U}{\sqrt{n \cdot \nu \cdot \hat{\epsilon}}}\]

API Reference

snntorch.weight_init.fluctuation_driven_normal_(tensor: Tensor, tau_mem: float, tau_syn: float, nu: float, mu_u: float = 0.0, sigma_u: float = 1.0, theta: float = 1.0, dt: float = 0.001, mode: str = 'analytical', generator: Generator | None = None) → Tensor[source]

Initialize weights from a normal distribution using fluctuation-driven initialization.

This method targets a membrane potential with mean mu_u and standard deviation sigma_u, placing neurons near threshold for optimal surrogate gradient learning.

Args:

tensor: Weight tensor to initialize (modified in-place). tau_mem: Membrane time constant in seconds. tau_syn: Synaptic time constant in seconds. nu: Estimated presynaptic firing rate in Hz. mu_u: Target mean membrane potential. Default: 0.0 (balanced). sigma_u: Target standard deviation of membrane potential. Default: 1.0. theta: Spike threshold. Default: 1.0. dt: Simulation timestep in seconds. Default: 1e-3. mode: “analytical” or “numerical” for PSP kernel integrals.

Default: “analytical”.

generator: Optional torch.Generator for reproducibility.

Returns:

The initialized tensor (same object, modified in-place).

Raises:

ValueError: If computed sigma_w is negative. ValueError: If tensor has fewer than 2 dimensions. ValueError: If mode is not “analytical” or “numerical”. ValueError: If any parameter is not positive.

Example:

>>> import torch.nn as nn
>>> import snntorch as snn
>>> fc = nn.Linear(784, 128)
>>> snn.weight_init.fluctuation_driven_normal_(
...     fc.weight, tau_mem=0.02, tau_syn=0.005, nu=10.0, sigma_u=1.0
... )

snntorch.weight_init.fluctuation_driven_normal_centered_(tensor: Tensor, tau_mem: float, tau_syn: float, nu: float, sigma_u: float = 1.0, dt: float = 0.001, mode: str = 'analytical', generator: Generator | None = None) → Tensor[source]

Initialize weights using centered fluctuation-driven initialization (mu_u=0).

This is a simplified version assuming balanced excitation/inhibition where the target mean membrane potential is zero. This is the most commonly used variant.

The weight distribution is N(0, sigma_w) where:: sigma_w = sigma_u / sqrt(n * nu * epsilon_hat)

(Eq. 9 from Rossbroich et al. 2022)

Args:

tensor: Weight tensor to initialize (modified in-place). tau_mem: Membrane time constant in seconds. tau_syn: Synaptic time constant in seconds. nu: Estimated presynaptic firing rate in Hz. sigma_u: Target standard deviation of membrane potential. Default: 1.0. dt: Simulation timestep in seconds. Default: 1e-3. mode: “analytical” or “numerical” for PSP kernel integrals.

Default: “analytical”.

generator: Optional torch.Generator for reproducibility.

Returns:

The initialized tensor (same object, modified in-place).

Raises:

ValueError: If tensor has fewer than 2 dimensions. ValueError: If mode is not “analytical” or “numerical”. ValueError: If any parameter is not positive.

Example:

>>> import torch.nn as nn
>>> import snntorch as snn
>>> conv = nn.Conv2d(64, 128, 3, padding=1)
>>> snn.weight_init.fluctuation_driven_normal_centered_(
...     conv.weight, tau_mem=0.02, tau_syn=0.005, nu=15.0
... )

snntorch.weight_init.compute_psp_kernel_integrals(tau_mem: float, tau_syn: float, dt: float = 0.001, mode: str = 'analytical') → Tuple[float, float][source]

Compute the integrals of the PSP kernel (epsilon_bar and epsilon_hat).

These integrals are used by the fluctuation-driven initialization to compute weight parameters from target membrane statistics.

Args:

tau_mem: Membrane time constant in seconds. tau_syn: Synaptic time constant in seconds. dt: Simulation timestep in seconds. Only used when mode=”numerical”. mode: “analytical” or “numerical”. Default: “analytical”.

Returns:

Tuple of (epsilon_bar, epsilon_hat).

Raises:

ValueError: If tau_mem or tau_syn are not positive. ValueError: If mode is not “analytical” or “numerical”.

Example:

>>> from snntorch import weight_init
>>> ebar, ehat = weight_init.compute_psp_kernel_integrals(
...     tau_mem=0.02, tau_syn=0.005
... )
>>> print(f"epsilon_bar={ebar:.4f}, epsilon_hat={ehat:.6f}")