primitive-lif-neuron

Primitive · LIFNeuron

Leaky Integrate-and-Fire neuron: a stateful, time-stepped non-linearity that converts continuous-valued inputs into binary {0, 1} spike trains over T timesteps. The spike-time dimension is added to the tensor shape: (B, N, D) → (B, T, N, D).

What it does

At each timestep t:

• Membrane potential: u[t] = beta * u[t-1] + x (leaky integration).
• Spike: s[t] = H(u[t] - theta) (Heaviside step at threshold).
• Soft reset: u[t] = u[t] - s[t] * theta (subtract threshold from neurons that fired).

Backward uses a surrogate gradient for the non-differentiable Heaviside — typically a triangular pulse centered at the threshold.

Interface (V1 stub)

from curry_train.primitives import LIFNeuron

lif = LIFNeuron(
    d_model=2048,
    T=4,                      # number of spike timesteps
    init_theta=1.0,           # firing threshold (learnable, per-feature)
    init_beta=0.5,            # decay before sigmoid (learnable, per-feature)
    gamma=1.0,                # surrogate gradient width
    backend="custom",         # or "spikingjelly"
)

# Input:  x  shape (B, N, D)
# Output: s  shape (B, T, N, D)  — binary spikes
spikes = lif(x)

Tensor shape contract

This is the only V1 primitive that changes tensor rank: it adds a T dimension.

• Pre-LIF: continuous (B, N, D).
• Post-LIF: binary (B, T, N, D).

Downstream layers (attention, MLP) must either:

• Operate on (B, T, N, D) directly (channel-aware spike processing), or
• Aggregate T first (e.g., spike rate mean over T → (B, N, D)) and then operate continuously.

The model's documentation should pin the rank contract at every boundary.

Backends

V1 should support two backends, selectable at construction:

• custom: pure-PyTorch implementation with torch.autograd.Function for the surrogate gradient. Slower, fully transparent, easy to debug.
• spikingjelly: SpikingJelly's ParametricLIFNode with optional CuPy backend. Faster, but requires pip install spikingjelly and CuPy for the fast path.

The reference implementation in csla_mt/model/spiking_neuron.py shows both.

Surrogate gradient choices

• Triangular (default): g(u) = max(0, 1 - |u - theta| / gamma) / gamma. Piecewise linear, centered at threshold.
• Sigmoid: g(u) = sigmoid(alpha * (u - theta)) * (1 - sigmoid(alpha * (u - theta))). Smooth.
• ATan: g(u) = (1/pi) * 1 / (1 + (alpha * (u - theta))²). Most common in modern SNN papers.

gamma (or alpha) controls the width: smaller = sharper, larger = smoother. Default 1.0 is a reasonable starting point; tune per task.

When to use

• Building or wrapping a spiking transformer (e.g., CSLA-MT clone).
• Energy-efficient inference targets (neuromorphic hardware compatibility).
• Exploring temporal-encoding variants.

When NOT to use

• Standard LLM / CV transformers — LIFNeuron is not a drop-in for ReLU/GELU; the time dimension changes everything downstream.

Boundaries

• Multi-step neurons require node.reset() between independent forward passes (for SpikingJelly). Forgetting this carries state across batches — silent bug.
• Mixed precision: SpikingJelly's CuPy backend only supports fp32/fp16; bf16 must be cast.
• The T dimension makes attention memory O(B * T * H * N²) — T=4 is a 4× memory tax on attention.

Implementation status

V1: stub at template/curry_train/primitives/lif_neuron.py. Reference implementation at /home/yanggl/code/autoresearch/csla_mt/model/spiking_neuron.py (the user's own working code).

Related

• skills/primitive-rmsnorm — for non-spiking layers; SNN often uses BatchNorm instead.
• The user's CSLA-MT code is the canonical reference; review that before reimplementing.
• Fang et al. 2021, "Incorporating Learnable Membrane Time Constant" (Parametric LIF basis).
• Neftci et al. 2019, "Surrogate Gradient Learning in Spiking Neural Networks".