NNT Define Linear Attention

Regular attention is quadratic: every token attends to every other token, so doubling the sequence length quadruples the compute. Linear attention is the family of tricks that breaks that scaling - approximate the softmax so the attention computation becomes linear in sequence length, and you can process way longer sequences on the same hardware. This node declares a linear-attention layer on the Neural Network Toolkit's layer stack.

The inputs

• embed_dim (default 512) and num_heads (default 8) - same contract as any attention layer; embed_dim should divide by num_heads.
• feature_map (default elu, also relu or softmax) - the kernel approximation. This is the heart of the trick. Softmax attention's similarity score gets replaced by a dot product of mapped features; elu (mapped as elu(x) + 1) and relu are the common choices because they're cheap and non-negative.
• eps (default 1e-6) - the numerical-stability epsilon in the denominator, guarding against division by near-zero normalization.
• causal (False) - whether to mask so tokens only attend to earlier tokens (autoregressive / decoder-style).
• dropout (0.1) and batch_first (True) - the usual suspects.

Output: LIST - the layer stack.

How it works

The idea in one line: instead of computing the full Q @ K^T attention matrix, rewrite attention as (Q' @ (K'^T @ V)) - note the grouping. By computing K'^T @ V first, you collapse the sequence dimension and the matrix multiply never has to materialize the N×N attention map. The toolkit's implementation does exactly this: project Q/K/V, apply the feature map, accumulate kv = K^T @ V, normalize with 1 / (Q @ sum(K) + eps), and project back out. The causal flag is recorded but - as with everything in this corner of the pack - the implementation that actually runs is a work in progress.

The honest caveat

Read the fine print for the pack's transformer nodes once and you've read it for all of them: the define nodes append layer dicts to the stack list, and the compile loop currently dispatches the classic layer types (Conv, Linear, Flatten, Pool, Norm). LinearAttention isn't in that dispatch list yet, even though a fully implemented LinearAttention module exists in the source - it just isn't wired into the model-build path. So today, this node is for learning the linear-attention parameterization and the stack format. If you need linear attention actually running on long sequences, grab a library that ships it - this is a teaching sandbox, and the author says exactly that in the README.

Common issues

• Layer missing from compiled model - expected in the current build; the definition travels, the instantiation doesn't (yet).
• Sequence length beyond training - even when it runs, long-sequence work tends to expose stability issues; that's what eps is there to babysit.

Install

Pack-level, nothing node-specific:

cd ComfyUI/custom_nodes
git clone https://github.com/inventorado/ComfyUI_NNT.git
cd ComfyUI_NNT
pip install -r requirements.txt

or ComfyUI Manager → "ComfyUI Neural Network Toolkit NNT", restart, and it's under NNT Neural Network Toolkit/Transformers.