Celsia's picture
Upload folder using huggingface_hub
ecc81b3 verified
|
Raw
History Blame Contribute Delete
8.98 kB

LTI and non-LTI mixers: what changes in N dimensions

The library's premise is that any 1-D mixer can be swept over an N-D lattice. That is true, and it is not uniform. Which mixer you sweep decides whether the sweep order is a modelling decision or a no-op, whether direction buys anything, and whether the whole stack collapses into a single N-D operator.

This document states those consequences and measures them. Every number here comes from tests/test_lti.py and td.testing.check_lti, and every claim that turned out to be folklore rather than fact is marked as such — there is one, and it is the interesting part.

The two properties

A mixer is a map from (M, A, H) to (M, A, H) — a batch of length-A sequences. Two properties matter:

Linear. f(x + y) = f(x) + f(y) and f(αx) = αf(x). Measured as relative error; a bias makes a mixer affine, which is reported separately because subtracting the constant recovers linearity.

Time-invariant. Delay the input and the output is delayed identically. Measured by zero-padding the front of the signal rather than by rolling it, because time invariance is a claim about a system at rest: feed it nothing, then feed the signal later, and the same thing should come out later. Rolling would wrap a different prefix into place and measure memory decay, which is a different question with a different answer.

td.testing.check_lti(lambda: td.mixers.ConvMixer(64, activation=None))
# ConvMixer: LTI (affine)
#     additivity          1.6e-16
#     homogeneity         1.4e-16
#     shift equivariance  0.0e+00
#     response to zero    3.1e-02

It never raises. No mixer is supposed to be LTI; this is a classification, not a check.

Every mixer this library ships

Measured at d_model=4, float64, eval mode. "Order gap" is the relative difference between sweeping h then w and sweeping w then h, with a shared mixer instance and no LayerNorm, so the only difference is order.

mixer linear time-invariant verdict order gap
scalar-valued filter (A ⊗ b) yes yes LTI 3e-16
ConvMixer(activation=None) yes yes LTI 3e-01
ConvMixer() (GELU) 4e-01 yes time-invariant, nonlinear 1e-01
TCNMixer 4e-01 yes time-invariant, nonlinear 4e-02
S4DMixer 5e-01 yes time-invariant, nonlinear 2e-02
S4Mixer 6e-01 yes time-invariant, nonlinear 4e-02
MambaMixer 5e-01 4e-04 neither 1e-02
LSTMMixer 6e-01 3e-03 neither 2e-02
GRUMixer 6e-01 3e-03 neither 5e-02
AttentionMixer 8e-02 2e-02 neither 2e-01

Read the first two rows together. They are the same class, both LTI, and their order gaps differ by fifteen orders of magnitude.

The correction: LTI does not mean order-free

The folklore — and my own first draft of this file — says: per-axis LTI operators commute, so for an LTI mixer the sweep order carries no information. The first half is true for scalar filters and the conclusion is false for every multichannel network anyone actually builds.

A 1-D convolution with groups=1 is a matrix-valued filter: offset j carries its own channel matrix W[:, :, j]. Composing along two axes gives

(F_h F_w x)[c] = Σ W[c,m,i] W[m,n,j] x[n, h+i, w+j]
(F_w F_h x)[c] = Σ W[c,m,j] W[m,n,i] x[n, h+i, w+j]

which agree only when W[:,:,i] W[:,:,j] = W[:,:,j] W[:,:,i] for every pair of offsets. Random matrices do not commute, so the order gap is 3e-01 — not a rounding artifact, a different model.

The operator commutes across axes exactly when it factorizes as a channel matrix times a spatial filter, W[out, in, j] = A[out, in] · b[j]. Then each axis contributes C_axis ⊗ A, the spatial parts commute because they act on different axes, and the channel parts are the same matrix. Constructed that way, the gap is 3e-16 — machine precision.

The depthwise row is the instructive one. A depthwise convolution is a scalar filter per channel and would commute; the pointwise mix that follows it is what breaks it. Separability across channels and commutation across axes are not the same property, and one does not imply the other.

So ScanPlan is not only for the selective models. A linear CNN swept over a lattice has an order-dependent result, and any implementation that reorders axes for convenience is changing the model.

What LTI does buy: exact separability

The property that survives is stronger and more useful than commutation.

For a linear mixer with no bias, a stack of one sweep per axis is exactly an N-D convolution whose kernel is the outer product of the per-axis kernels — contracted in the order the sweeps happened. Not an approximation:

tests/test_conv.py::test_separable_stack_equals_a_full_nd_convolution
    rank 2 and rank 3, checked against F.conv2d / F.conv3d, max diff < 1e-12

This is the convolutional twin of the kernel family's Kronecker identity, and it is what "S4ND applies its axes simultaneously" rests on: with an LTI mixer the sequential sweeps and the joint kernel are the same operator, so a single FFT over all axes is available. The joint kernel still depends on the order — separability and order-independence are different claims, and only the first one holds.

Checked against S4ND itself, not just against the theory. S4ND does not sweep: it builds one 1-D kernel per axis, outer-products them in Fourier space, and applies a single N-D FFT convolution. Transcribing that composition as an oracle and feeding it our own kernels:

tests/test_published_composition.py
    our sequential per-axis sweep vs S4ND's simultaneous N-D kernel
    rank 2 and rank 3, max relative difference 1.8e-17

So a plan of one-axis-per-layer sweeps reproduces, exactly, a model that was never written as a sweep. That is the strongest single piece of evidence for the library's premise, and it is available because S4ND's kernels are diagonal in channels — the scalar-filter case above. The negative control in that file gives the axes a channel-mixing filter and the identity dies, which is the same correction stated in the other direction.

Two things follow for free:

  • Parameter cost. A separable stack is r·k parameters where a full N-D kernel is k^r. That is the whole bargain, and it is why the rank-1 restriction is worth stating: a separable stack cannot represent a diagonal edge detector, and depth plus channel mixing is what buys most of it back.
  • Direction is nearly free. For a centred convolution, sweeping backwards is the same operator with a mirrored kernel. Constructed symmetric, forward and reverse agree to 1e-12; an LSTM disagrees by 1e-3 or more. So bidirectional= earns its keep for causal and recurrent mixers and buys a mirrored filter for a CNN.

Why the non-LTI models are the way they are

Recurrences are not time-invariant, despite being perfectly causal. A gated recurrence fed zeros does not sit still: its biases drive the state, so a delayed signal arrives to a machine in a different configuration than the original did. The deviation is small (3e-03 for nn.LSTM) because the state settles — it is a transient, not a wholesale failure — but it is real, and it is the reason "causal" and "time-invariant" are separate columns in the table.

Attention is permutation-equivariant, which is stronger than shift-equivariant on a fixed token set, and still fails this test. Delaying a signal means padding it, and padding adds tokens to attend over. The token set changes, so the output does. This is not a defect; it is the reason positional embeddings exist at all.

Mamba's selectivity is the point. Its A, B, C, dt are functions of the input, so there is no fixed operator to commute with anything, no joint kernel to build, and no simultaneous application. Every axis must be swept in sequence, the order matters, and the direction matters — which is exactly why Mamba-ND needed a schedule and S4ND did not. ScanPlan.paired() exists because of the 1e-02 in that row.

Practical consequences

if your mixer is… then…
a scalar-valued LTI filter sweep order is genuinely irrelevant; pick any plan
LTI, multichannel the stack is one N-D kernel, but the order changes which one
time-invariant, nonlinear (S4, TCN, conv) order matters; direction matters for causal ones
neither (Mamba, RNN, attention) order and direction are architectural choices — measure them

The library does not enforce any of this. It measures it, and it gives the schedule a name so the choice is visible in a config file instead of buried in a loop.

Reproducing

pytest tests/test_lti.py -q
import torch_dimensions as td

print(td.testing.check_lti(lambda: td.mixers.MambaMixer(64)))