Skip to content

Fix LeakyParallel silently ignoring a per-neuron beta; correct the reset=zero formula in Leaky - #441

Merged
ixfd64 merged 1 commit into
jeshraghian:masterfrom
tritsystem:fix/leakyparallel-per-neuron-beta
Sep 20, 2026
Merged

ixfd64 merged 1 commit into
jeshraghian:masterfrom
tritsystem:fix/leakyparallel-per-neuron-beta

Conversation

@tritsystem

@tritsystem tritsystem commented Sep 1, 2026

Copy link
Copy Markdown
Contributor

Fixes #442.

Problem

1. LeakyParallel silently drops a per-neuron beta.

_beta_to_weight_hh walks an if / elif / elif / else chain:

if isinstance(self.beta, float) or isinstance(self.beta, int):
    self.rnn.weight_hh_l0.fill_(self.beta)
elif isinstance(self.beta, torch.Tensor) or isinstance(self.beta, torch.FloatTensor):
    if len(self.beta) == 1:
        self.rnn.weight_hh_l0.fill_(self.beta[0])
elif len(self.beta) == self.hidden_size:          # sibling of the Tensor elif
    for i in range(self.hidden_size):
        self.rnn.weight_hh_l0.data[i].fill_(self.beta[i])
else:
    raise ValueError(...)

_beta_buffer always stores beta as a torch.Tensor, so control always
enters the second branch, and only its inner if len(self.beta) == 1 is
reachable. The elif len(self.beta) == self.hidden_size (the per-neuron path)
and the else: raise ValueError are dead.

Passing beta=<tensor of length hidden_size> — which the docstring explicitly
supports ("multi-valued (one weight per neuron)") — leaves
rnn.weight_hh_l0 at its random RNN initialisation, with no error. The
layer then trains and runs with arbitrary recurrent decay rates. A
wrong-length beta is likewise accepted silently.

import torch, snntorch as snn
lp = snn.LeakyParallel(input_size=4, hidden_size=6, beta=torch.linspace(0.1, 0.9, 6))
torch.diagonal(lp.rnn.weight_hh_l0)
# tensor([ 0.131, 0.136, -0.121, -0.028, 0.308, 0.287])   <- random init, not the betas
snn.LeakyParallel(input_size=4, hidden_size=6, beta=torch.tensor([0.3, 0.4]))
# no error

Reproduce (before this PR):

reproduce

2. Leaky docstring — reset_mechanism="zero" formula is wrong.

It states U[t+1] = βU[t] + I_syn[t+1] - R(βU[t] + I_in[t+1]), i.e.
(1-R)(βU[t] + I), which is 0 on the step after a spike regardless of
input
. The implementation (_base_zero) is standard reset-then-integrate,
U[t+1] = β(1-R)U[t] + I_in[t+1], which equals I_in on that step (verified
against a hand-rolled reference — matches to 0.0). Leaky also has no
synaptic current, so I_syn is a copy-paste from Synaptic.

Changes

  • snntorch/_neurons/leakyparallel.py — nest the two length checks and the
    ValueError inside the torch.Tensor branch so the per-neuron path is
    reached; an unsupported beta type now raises TypeError instead of
    falling through.
  • snntorch/_neurons/leaky.py — fix the reset_mechanism="zero" formula and
    drop the stray I_syn.

Tests

New tests/test_snntorch/test_leakyparallel.py (5 tests). Three fail on
master without this change:

test on master
scalar beta fills the diagonal pass
per-neuron beta (len == hidden_size) written to the diagonal fail
length-1 beta tensor still works pass
wrong-length beta raises ValueError fail (silently accepted)
two layers identical but for a per-neuron beta give different output fail (both ran at default recurrent weights)

Full suite: 195 → 200 passed, 2 xfailed.

After this PR:

after fix

Checklist

  • Applied flake8 and black (changed files clean; flake8 --select=E9,F63,F7,F82 clean)
  • Implemented unit tests and they all pass

…set="zero" formula in Leaky

### Problem

**1. `LeakyParallel` drops a per-neuron `beta`.**
`_beta_to_weight_hh` walks an `if / elif / elif / else` chain:

    if isinstance(self.beta, (float, int)): ...
    elif isinstance(self.beta, torch.Tensor) or isinstance(self.beta, torch.FloatTensor):
        if len(self.beta) == 1: ...
    elif len(self.beta) == self.hidden_size:   # <-- sibling of the Tensor elif
        for i in range(self.hidden_size): ...
    else:
        raise ValueError(...)

`_beta_buffer` always stores `beta` as a `torch.Tensor`, so control always
enters the second branch and only its inner `if len(self.beta) == 1` is
reachable. The `elif len(self.beta) == self.hidden_size` (the per-neuron
path) and the `else: raise ValueError` are dead. Passing
`beta=<tensor of length hidden_size>` -- which the docstring explicitly
supports ("multi-valued (one weight per neuron)") -- leaves
`rnn.weight_hh_l0` at its random RNN initialization, with no error. The
layer then trains and runs with arbitrary recurrent decay rates instead
of the requested ones.

**2. `Leaky` docstring for `reset_mechanism="zero"` is wrong.**
It states `U[t+1] = βU[t] + I_syn[t+1] - R(βU[t] + I_in[t+1])`, i.e.
`(1-R)(βU[t] + I)`, which is `0` on the step after a spike regardless of
input. The implementation (`_base_zero`) is standard reset-then-integrate,
`U[t+1] = β(1-R)U[t] + I_in[t+1]`, which equals `I_in` on that step. Also
`Leaky` has no synaptic current, so `I_syn` is a copy-paste from
`Synaptic`.

### Changes

- `snntorch/_neurons/leakyparallel.py`: nest the two length checks (and
  the `ValueError`) inside the `torch.Tensor` branch so the per-neuron
  path is reached; a genuinely unsupported `beta` type now raises
  `TypeError` instead of falling through silently.
- `snntorch/_neurons/leaky.py`: fix the `reset_mechanism="zero"` formula
  and drop the stray `I_syn`.

### Tests

New `tests/test_snntorch/test_leakyparallel.py` (5 tests):

- scalar `beta` still fills the diagonal,
- **per-neuron `beta` (len == hidden_size) is written to the diagonal**
  (fails without this change),
- length-1 `beta` tensor still works,
- **a bad-length `beta` raises `ValueError`** (fails without this change --
  it was silently accepted),
- **two layers identical but for a per-neuron `beta` produce different
  outputs** (fails without this change -- both ran at the default
  recurrent weights).

Full suite: 195 -> 200 passed, 2 xfailed.
@ixfd64

ixfd64 commented Sep 4, 2026

Copy link
Copy Markdown
Collaborator

I don't see any issues. Let me know once the PR is ready for review.

@tritsystem
tritsystem marked this pull request as ready for review September 4, 2026 07:13
@tritsystem

Copy link
Copy Markdown
Contributor Author

Thanks for taking a look @ixfd64 — I've just moved it out of draft, it's ready for review now.

Quick recap of what it fixes:

  1. LeakyParallel silently ignores a per-neuron beta. _beta_buffer always stores beta as a tensor, so in _beta_to_weight_hh the elif len(self.beta) == self.hidden_size branch and the else: raise are unreachable siblings of the isinstance(..., torch.Tensor) elif. Passing beta of length hidden_size (which the docstring says is supported) leaves rnn.weight_hh_l0 at its random RNN init, no error. The fix nests the two length checks + the error inside the tensor branch.
  2. Leaky docstring reset_mechanism="zero" formula describes (1-R)(βU + I) (→ 0 the step after a spike regardless of input); the implementation is standard reset-then-integrate (→ I_in on that step). Also drops a stray I_syn copy-pasted from Synaptic.

5 new tests in test_leakyparallel.py, 3 fail on master; full suite 195 → 200 passed. CI green (3.9/3.10/3.11). Linked issue #442 has the standalone write-up.

@ixfd64
ixfd64 merged commit de0a276 into jeshraghian:master Sep 20, 2026
3 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

LeakyParallel silently ignores a per-neuron beta (tensor of length hidden_size)

2 participants