exact-event-adjoint

Grid-free exact continuous-time adjoint gradients for Time-To-First-Spike (TTFS) spiking neural networks. No surrogate gradients. No dense time grid. Every gradient is the analytic derivative of the threshold-crossing event times, and every claim in this file is backed by a numerical test in tests/.

phi_k(t) = sum_m W[o,m] * K(t, t_in[m])  -  (theta - u_reset) * sum_{j<k} R(t, t_j)  -  theta  = 0

• K(t, s) = (tau_s*tau_m)/(tau_m-tau_s) * (exp(-(t-s)/tau_m) - exp(-(t-s)/tau_s)) double-exponential postsynaptic kernel.
• R(t, s) = exp(-(t-s)/tau_m) decaying post-spike reset kernel (subtractive reset).

Installed features (each one was verified empirically in tests/)

Precise wording (we do not overclaim)

These are the exact, defensible statements you may make in a paper:

1. Continuous-Time IFT & Saltation resetting — TRUE. The first firing time is an implicit function of the latent variables and its gradient is the exact IFT derivative — no surrogate approximation is used. For multi-spike trains the implicit equations for all spike times are coupled through the reset kernels; the backward pass solves that coupled system exactly. The coupling terms b_{k,j} are precisely the saltation-type reset multipliers -(theta-u_reset)/tau_m * R(t_k, t_j) of the threshold-reset discontinuity: the classic instantaneous linearization 1 + (theta-u_reset)/(tau_m u_dot) is the degenerate short-gap limit, ours is the exact finite-gap generalization.

2. Priority-queue, grid-free — TRUE, with one precision note. The package evaluates all membrane quantities on event-ordered sets (presynaptic times, kernel-peak lags, previous spikes, t_max) and never builds a uniform time grid. Probe counts scale with the number of events, not with the horizon length or any step size. The label "O(1) memory" that circulated in earlier demos is not retained: the heap-based engine is O(number of pending events) memory with O(log) work per push. That is what is actually true, so that is what we ship.

3. Escape-Noise Regularization — TRUE with a name correction. The implemented regularizer is a deterministic peak-margin penalty (max of the differentiable soft-weighted probe peak below threshold) that provably removes dead/silent-neuron gradient collapse: gradient descent raises silent weights. This is not the stochastic escape-rate functional; we therefore call it the existence regularization (ExactAdjointFunction(use_existence=True) or existence_loss_and_grads).

4. Full Layer Suite — TRUE. ExactAdjointLinear, ExactAdjointConv2d, ExactAdjointMultiSpike, ExactAdjointRecurrent / ExactAdjointTTFSRnn, ExactAdjointAttention, ExactAdjointMultiHeadAttention (genuine per-head split, sqrt(head_dim) scaling), ExactAdjointFFN, with AdjointSpikeNorm, latency / spike-count losses. All constructed layers run forward and backward.

Quickstart

import torch
from exact_event_adjoint import (
    ExactAdjointLinear, ExactAdjointMultiSpike,
    latency_cross_entropy, AdjointSpikeNorm,
)

torch.manual_seed(0)
encoder = ExactAdjointLinear(4, 6, tau_m=15.0, tau_s=4.0, theta=1.0, t_max=50.0)
decoder = ExactAdjointMultiSpike(6, 3, max_spikes=4)
norm    = AdjointSpikeNorm(6)

t_in = torch.rand(8, 4) * 40 + 1          # latency-encoded input spikes
t1   = encoder(t_in)                      # (B, O) exact first-spike times
h    = norm(t1)
train = decoder(h)                        # (O, B, K) reset-aware spike trains

loss = latency_cross_entropy(t1, torch.randint(0, 6, (8,)))
loss.backward()                           # exact IFT gradients, no surrogates

Mixed firing/silent training uses the existence flag:

layer = ExactAdjointLinear(4, 3, use_existence=True)   # dead neurons get gradient pressure

The exact per-layer alternative with the same maths lives in ExactEventLinear / ExactEventFunction (bit-identical to ExactAdjointFunction).

Running the verification suite

python -m pytest exact_event_adjoint/tests -q
# 20+ tests: finite-difference Jacobians, brute-force simulations, engine-vs-solver,
# existence direction, head split, layer suite, grid-free static guard

Integrity notes

• The solver (solver.py) performs no autograd; all broadcast arithmetic is ordinary tensor ops and every gradient is provided analytically by the autograd Function backward passes.
• Recurrent loops connect spike-latency outputs to latency-encoded inputs (w_in + w_rec); this is differentiable TTFS feedback, not a continuous dynamical-system solver for the recurrent loop (scope note).
• EventDrivenO1Engine.simulate_events implements feedforward receptive fields (static inputs). Causal recurrent rescheduling is future work.

License

MIT — see LICENSE.