Nodes/ComfyUI Neural Network Toolkit NNT /NNT Define Vanilla Attention
ComfyUI Node

NNT Define Vanilla Attention

Five attention recipes in one node

By inventorado·Created 2 years ago·Updated 2 years ago· 69
NNT Define Vanilla Attention
  • LAYER_STACK
  • LIST
embed_dim512
attention_typescaled_dot_product
dropout0.1
use_biasTrue
add_zero_attnFalse
batch_firstTrue

Every modern generative model leans on attention, and at some point you'll want to see the mechanism in isolation instead of buried inside a transformer block. NNT Define Vanilla Attention is that node: it adds a standalone attention layer to your NNT stack, with five different attention flavors behind a single attention_type dropdown. It's the "anatomy lesson" node of the transformer family in this pack.

What it actually does

Like the other define nodes in inventorado/ComfyUI_NNT, it appends a layer description to the LAYER_STACK list that NntCompileModel turns into a compiled model. Attention in its plainest form answers "given a query, how much should I look at each other token?" - it weights other positions by similarity and sums them up. The node wraps that idea in a few concrete formulas:

  • scaled_dot_product - query·key similarity, divided by the square root of the embedding dimension so the values don't explode, then softmax. This is the modern default, used by basically every transformer.
  • dot_product - the same thing without the scaling; simpler, and numerically worse for large embeddings.
  • additive - the older Bahdanau-style attention that learns a small network to score pairs. Predates scaled dot product and mostly appears in history sections now.
  • relative - injects relative distance between tokens into the scores; this is the idea that made Transformer-XL and friends work well.
  • local - only attends within a window around each position. Cheap, and a stepping stone to the Reformer-style tricks.

Inputs that matter

  • embed_dim - the width of the input (and output) representation. 512 is the default and a fine starting point.
  • attention_type - pick the recipe above. scaled_dot_product is the one you'll actually want 95% of the time.
  • dropout - applied to the attention weights after softmax; 0.1 is the standard starting value.
  • use_bias - whether the projection layers get biases. Usually fine to leave True.
  • add_zero_attn - appends an all-zero key/value pair so the layer can "attend to nothing." Rarely needed.
  • batch_first - keep True for the friendly [batch, seq, embed_dim] layout.

The honest gotcha

The recurring caveat applies here more than anywhere: this pack is a self-described work in progress, and the attention nodes are its most experimental corner. The definitions are clean and instructive, but the current compile path doesn't guarantee every attention entry materializes in the compiled model yet. Use NntCompileModel's "Only create script" mode to see what actually gets built, and don't build a production pipeline on these nodes - build understanding on them. If you want to feel the difference between scaled dot product and additive attention, that's exactly what this node is for.

Installing NNT

Part of inventorado/ComfyUI_NNT. ComfyUI Manager (search "ComfyUI Neural Network Toolkit") or:

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

Restart ComfyUI after. The dependency list is a heavy scientific stack - torch, numpy, scikit-learn, pandas, transformers, shap - so the first install takes a while. The pack's example workflows want ComfyUI-Jjk-Nodes for text output; let Manager fetch it.

CategoryNNT Neural Network Toolkit/Transformers

Inputs (7)

NameTypeDefaultDescription
embed_dimINT51264–2048
attention_typeCOMBOscaled_dot_product5 options: dot_product, additive, scaled_dot_product, relative, local
dropoutFLOAT0.10–0.9
use_biasCOMBOTrue2 options: True, False
add_zero_attnCOMBOFalse2 options: True, False
batch_firstCOMBOTrue2 options: True, False
LAYER_STACKoptLIST

Outputs (1)

NameTypeDescription
LISTLIST