SnAp: sparse n-step recurrent influence

SnAp: sparse n-step recurrent influence#

SnAp-n retains recurrent influence inside an \(n\)-step dependency neighborhood. A dense recurrent matrix becomes all-to-all almost immediately, so this tutorial uses a sparse ring where increasing \(n\) has a visible structural meaning.

1. Principle#

Let \(A\) be the directed dependency graph of recurrent hidden positions. SnAp-n keeps trace blocks reachable within the requested order and drops more distant influence. As \(n\) grows, the neighborhoods nest; once every position is reachable, the within-group trace saturates. Memory grows with the retained neighborhood width, so \(n\) is an accuracy-cost coordinate rather than a generic optimizer setting.

2. Applicability and exclusions#

Use SnAp when recurrence is structurally sparse and its position graph can be derived by the compiler.

Do not use a dense recurrence to demonstrate scaling: its graph diameter is one, so small \(n\) already saturates. Unsupported or position-relabeling tails may trigger conservative widening or explicit rejection; inspect diagnostics before interpreting an estimate.

import brainevent
import brainstate
import braintrace
import jax.numpy as jnp

brainstate.random.seed(41)


def ring_structure(n_rec):
    row = jnp.arange(n_rec)
    mask = jnp.zeros((n_rec, n_rec))
    mask = mask.at[row, row].set(1.0)
    mask = mask.at[row, (row + 1) % n_rec].set(1.0)
    return brainevent.CSR.fromdense(mask), int(mask.sum())


class SparseRingRNN(brainstate.nn.Module):
    def __init__(self, n_in=2, n_rec=7):
        super().__init__()
        self.csr, nnz = ring_structure(n_rec)
        self.w = brainstate.ParamState(0.25 * brainstate.random.randn(nnz))
        self.w_in = brainstate.ParamState(
            0.2 * brainstate.random.randn(n_in, n_rec)
        )
        self.h = brainstate.HiddenState(jnp.zeros(n_rec))

    def update(self, x):
        recurrent = braintrace.sparse_matmul(
            self.h.value, self.w.value, sparse_mat=self.csr
        )
        self.h.value = jnp.tanh(x @ self.w_in.value + recurrent)
        return self.h.value

3. Inspect neighborhood growth#

Each learner receives a fresh model with the same seeded parameter values. The reported K is the compiled neighborhood width, not an inferred value from the requested order.

x0 = brainstate.random.randn(2)


def compiled_width(order):
    with brainstate.random.seed_context(43):
        model = SparseRingRNN()
    learner = braintrace.compile(model, braintrace.SnAp, x0, n=order)
    group = learner.graph.hidden_groups[0]
    if group.snap is None:
        return 1, group.trace_state_width, False
    return group.snap.num_neighbour, group.trace_state_width, group.snap.is_saturated


snap1 = compiled_width(1)
snap2 = compiled_width(2)
snap3 = compiled_width(3)

print("SnAp-1 (K, trace width, saturated):", snap1)
print("SnAp-2 (K, trace width, saturated):", snap2)
print("SnAp-3 (K, trace width, saturated):", snap3)
print("nested widths:", snap1[0] <= snap2[0] <= snap3[0])

SnAp-1 (K, trace width, saturated): (1, 1, False)
SnAp-2 (K, trace width, saturated): (2, 2, False)
SnAp-3 (K, trace width, saturated): (3, 3, False)
nested widths: True

4. Execute the public sequence path#

The next cell checks that a non-saturated order produces a finite online gradient. It does not compare that approximation with BPTT; a finite-window oracle is required for an accuracy claim.

with brainstate.random.seed_context(43):
    model = SparseRingRNN()
inputs = brainstate.random.randn(8, 2)
targets = jnp.zeros((8, 7))
learner = braintrace.compile(model, braintrace.SnAp, inputs[0], n=2)


def step_loss(x, target):
    return jnp.mean((learner(x) - target) ** 2)


brainstate.nn.reset_all_states(model)
learner.reset_state()
grads, loss = learner.etrace_grad(
    inputs,
    targets,
    step_fn=step_loss,
    loss_output="scalar",
    return_value=True,
)
gradient_norm = jnp.sqrt(sum(jnp.sum(g * g) for g in grads.values()))

print(f"sequence loss: {float(loss):.4f}")
print("finite nonzero gradient:", bool(jnp.isfinite(gradient_norm) & (gradient_norm > 0)))

sequence loss: 0.0816
finite nonzero gradient: True

5. Interpretation and limits#

The ring shows that the compiled neighborhood grows with order instead of saturating immediately. This is structural evidence that the example exercises SnAp’s defining mechanism. It is not evidence that order 2 is adequate for a particular task. Accuracy, compiler conservatism, and memory must be evaluated together on the intended sparse architecture.

Reference and API#