Multi-Head Attention

You already know dot-product attention - that's the one attention. Multi-head attention is the trick that runs it several times in parallel on different learned projections, then glues the results back together, and it's the version every real transformer block uses. ComfyDL's CdlMultiHeadAttention builds that whole layer from your parameters, and it's the node you'll reach for when you're hand-assembling a Transformer encoder block instead of using a one-shot encoder node.

How it works

The idea: one attention head is a single view of the sequence - it can learn to look for, say, the subject-verb relationship, but it can't easily track several different relationships at once. Multi-head fixes that by giving each head its own projection.

Concretely, the layer has four learned linear projections: W_q, W_k, W_v for the inputs, and W_o to mix heads back together. Queries, keys, and values are each projected to num_hiddens width, then reshaped into num_heads parallel slices of width num_hiddens / num_heads instead of using num_heads separate networks. Each slice runs scaled dot-product attention independently, the heads are concatenated back into full width, and W_o projects the result. That reshape trick is the "multi" - one wider layer doing the work of several narrow ones, which is also why it's fast on a GPU.

Because it's an nn.Module, it exposes the per-head attention weights after a forward pass. Visualize those and you can literally see heads specialising - one tracking position, one tracking syntax. That demo is half the fun of this node.

Inputs and outputs that matter

• num_hiddens (default 8) - the model width. This is the dimension of your embeddings/features, and it must be divisible by num_heads, because each head gets num_hiddens / num_heads channels.
• num_heads (default 4) - how many parallel heads to run. Powers of two are conventional; 8 heads at width 512 is the transformer-classic.
• dropout (default 0) - on the attention weights inside each head.
• use_bias (default false) - whether the four projections get bias terms. Off matches the original "Attention Is All You Need" design; flipping it on is harmless for experiments.

The model output (a cdlModel) expects batch-first (batch_size, seq_len, num_hiddens) queries, keys, and values, plus an optional valid_lens, and returns (batch, num_queries, num_hiddens). Wire it into a CdlModelForward node to test it, or pair it with CdlAddNorm and CdlPositionWiseFFN to assemble a Transformer encoder block by hand.

Installing ComfyDL

cd ComfyUI/custom_nodes
git clone https://github.com/Cynthia-lxx/ComfyDL ./ComfyDL
pip install -r ./ComfyDL/requirements.txt

Restart after that. ComfyDL's deps are light (matplotlib, IPython, matplotlib-inline) - no model files here. Manager: search "ComfyDL", but this young educational pack may not be in the catalog yet, so the clone route is the reliable one.

Gotchas

The one error people actually hit is a num_hiddens that isn't divisible by num_heads - the node will refuse and tell you, because 8 hidden units can't split into 3 heads. So pick your width and head count together: width 8 works with 1, 2, 4, or 8 heads, not 3. And remember the width you set here has to match the embedding width of whatever tensors you feed it later - set num_hiddens to your embedding size, not to a number that feels lucky.