NNT Define Multihead Attention

Multi-head attention is the layer that made transformers transformers - split the input into heads, let each head attend to different relationships, concatenate the results. This node is the Neural Network Toolkit's way of declaring one on a layer stack. The inputs map almost exactly onto PyTorch's nn.MultiheadAttention, so if you know one, you know the other.

The inputs

• embed_dim (default 512) - the embedding dimension; it must be divisible by num_heads.
• num_heads (default 8) - attention heads. The classic rule of thumb is embed_dim / num_heads = 64 per head, so 512 with 8 heads is the textbook setup.
• dropout (default 0.1) - applied on attention weights during training.
• bias (True) - bias in the projection layers.
• add_bias_kv (False) - an extra learned bias on the key/value projections, a niche trick used in some encoder-decoder setups.
• add_zero_attn (False) - a dummy "zero" position at the end of the key/value sequences; occasionally helps in sequence alignment.
• batch_first (True) - input layout (batch, seq, embed). The toolkit defaults to True so you don't fight PyTorch's historical (seq, batch, embed) default.
• kdim / vdim (0) - key/value dimensions when they differ from embed_dim (cross-attention where the encoder/decoder widths differ). 0 means "same as embed_dim."

Output: LIST (the layer stack, chain it onward).

How it works - and where the pack stands

The define node appends {'type': 'MultiheadAttention', 'embed_dim': 512, ...} to the stack. NntCompileModel is the stage where stack entries become real nn.Modules.

Honest status check from the source: the compile loop currently dispatches the classic layer types - Conv, Linear, Flatten, Reshape, pooling, norms. Multi-head attention appears in the activation menu of the compile node (you can use nn.MultiheadAttention as an activation function there) and there's attention-building machinery in the codebase, but a stack entry of type MultiheadAttention isn't wired into the model-build path yet. The pack is a work-in-progress teaching tool - the author's README says as much - so the practical advice is: use this node to learn what an attention definition looks like and how the stack format works, and don't be surprised if the compiled model doesn't contain it today.

If you want to see multi-head attention actually run, the more reliable route inside the pack is the activation_function menu on NntCompileModel, which can insert a real nn.MultiheadAttention into the model.

Common issues

• embed_dim not divisible by num_heads - PyTorch will reject it at build time. Keep 512/8 unless you have a reason.
• Layer missing from compiled model - expected in the current build; that's the WIP gap, not your wiring.

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, and find it under NNT Neural Network Toolkit/Transformers.