NNT Define Alibi PositionalBias

Attention with a Linear Bias (ALiBi) is the elegant alternative to learned positional encodings: instead of adding position vectors to the embeddings, you just subtract a distance penalty from the attention scores themselves. Closer tokens score higher, far tokens get docked - no extra parameters, and it generalizes to longer sequences than it was trained on. It's the trick behind several long-context models, and this node declares it on the Neural Network Toolkit's layer stack.

The inputs

• num_heads (default 8) - each head gets a different slope, so this sets how many distance penalties you get.
• max_seq_length (default 512) - the sequence length the bias table is precomputed for.
• causal (False) - mask so tokens only attend to the past (decoder style).
• slope_multiplier (default 1.0) - scales all head slopes up or down. This is the one to poke if you want to experiment: higher multiplier means attention falls off with distance faster.

Output: LIST - the layer stack.

How it works

The math is refreshingly simple. The toolkit computes per-head slopes using the standard ALiBi formula - head i gets slope 2^(-8i/num_heads) (with the standard fallback for non-power-of-two head counts). Then it builds a distance matrix: position j attending to position k gets penalized by slope * |j - k|. That bias gets added straight onto the attention scores before softmax. If causal is on, the upper triangle is masked to -inf as well. The module's forward just does attention_scores + alibi_bias, sliced to the current sequence length.

That's the design, clean as can be. The status check, though, is the same one that runs through this pack's transformer nodes: the define nodes append layer dicts to the stack list, and the compile loop currently instantiates the classic layer types - Conv, Linear, Flatten, Reshape, pooling, norms. An AlibiPositionalBias stack entry isn't dispatched into the built model yet. The full AliBi module with slope computation exists in the source, but the compile path doesn't call it. So today, this node is a learning aid for ALiBi's parameterization, not a component that changes your compiled model.

Common issues

• Layer missing from compiled model - expected in the current build; the WIP gap, not your wiring.
• max_seq_length too small - if/when it runs, sequences longer than the precomputed table get truncated to it, so set it to your intended ceiling.

Install

Pack-level:

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.