NNT DefineLocal Attention

Full attention lets every token peek at every other token, which is powerful and expensive. Local attention says: most of that peeking is a waste, so give each token a window and let it attend only to its neighbors. It's one of the classic tricks for scaling transformers to long sequences, and it's what this node declares on the Neural Network Toolkit's layer stack.

The inputs

• embed_dim (default 512), num_heads (default 8) - the usual attention contract.
• window_size (default 128) - the maximum distance over which attention is allowed. This is the headline parameter; bigger window = more context but more compute.
• look_behind (default 64) and look_ahead (default 0) - fine-grained control over the window shape. look_ahead=0 with a positive look_behind gives you a causal-ish window (only the past), which is the decoder pattern. Set both positive for a symmetric band around each token.
• autopad (True) - pad short sequences so the windowing math stays clean near the edges.
• dropout (0.1), batch_first (True) - the usual suspects.

Output: LIST - the growing layer stack.

How it works

The toolkit's implementation builds a boolean attention mask where row i is allowed to attend to positions [i - look_behind, i + look_ahead] and masks everything else to -inf before softmax. So instead of the N×N attention matrix being dense, it's banded - a diagonal stripe of allowed positions. The window_size sets the overall reach; look_behind/look_ahead shape the stripe.

That's the theory. Now the status check that applies to all of this pack's transformer nodes: the define nodes write layer dicts onto the stack list, and the compile loop currently instantiates the classic layer types - Conv, Linear, Flatten, Reshape, pooling, norms. A LocalAttention stack entry isn't in that dispatch list yet. A fully coded LocalAttention module with the masking logic exists in the source, but the compile path doesn't call it. So treat this as a node for studying local attention's parameters and the stack format, and don't expect the compiled model to contain the layer yet.

Common issues

• Layer missing from compiled model - expected in the current build; it's the WIP gap shared across the pack's attention nodes, not your wiring.
• Window vs. context - remember the tradeoff you're declaring: a small window means the model literally cannot see long-range dependencies. That's the whole game with local attention - you accept it to save compute.

Install

Pack-level install:

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, then look under NNT Neural Network Toolkit/Transformers.