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

NNT Define TransformerXL Attention

The node that gave transformers a memory

By inventorado·Created 2 years ago·Updated 2 years ago· 69
NNT Define TransformerXL Attention
  • LAYER_STACK
  • LIST
d_model512
num_heads8
mem_len512
same_lengthFalse
clamp_len-1
dropout0.1
batch_firstTrue

Transformers are great at context, but only as far as the window you feed them - and training on long windows is expensive. NNT Define TransformerXL Attention is the node that teaches you the trick that broke that limit: a recurrence-like segment-level memory that carries hidden states from previous chunks into the next one, so the model effectively sees further back than its input window. It's a genuinely important idea in the history of sequence models, and this is one of the few places you can poke at it with a mouse.

What it actually does

It appends a TransformerXLAttention entry to the LAYER_STACK list that NntCompileModel compiles into a model. The defining feature is mem_len: when processing a sequence in segments, the layer keeps the last mem_len hidden states from earlier segments and lets attention look at them alongside the current segment. That gives the model a soft long-term memory without retraining on longer inputs. Transformer-XL also introduced a relative position scheme (distance-based rather than absolute), which is why you'll often see it paired with the relative-position-bias node in this pack.

Inputs that matter

  • d_model - embedding width; 512 default.
  • num_heads - attention heads.
  • mem_len - the memory length: how many previous hidden states get carried forward. 0 disables the memory entirely (making this a plain attention layer - handy for A/B testing what the memory actually buys you). The default 512 is a real memory cost, so don't be shy about dialing it down.
  • same_length - when True, positions near the start of a segment attend to the same number of past positions as everyone else (used to make certain kinds of evaluation fair).
  • clamp_len - caps how far back relative positions are allowed to extend; -1 means "no clamp." Clamping helps keep the relative-position embeddings from being over-trained on tiny ranges.
  • dropout, batch_first - the usual; keep batch_first on.

The honest gotcha

The recurring caveat, louder here: this pack is a self-described work in progress and the transformer section is its most experimental territory. The node produces a clean definition that's great for understanding what segment memory is, but the current compile path doesn't guarantee every transformer entry materializes in the compiled model yet. Check what actually builds with NntCompileModel's "Only create script" mode before investing in a big architecture. This is a teaching node, not a FlashAttention replacement - its job is to make the memory idea click.

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 requirements are a heavy scientific stack - torch, scikit-learn, pandas, transformers, shap - so the first install takes a while. The pack's example workflows also want ComfyUI-Jjk-Nodes for their text displays.

CategoryNNT Neural Network Toolkit/Transformers

Inputs (8)

NameTypeDefaultDescription
d_modelINT51264–2048
num_headsINT81–32
mem_lenINT5120–2048
same_lengthCOMBOFalse2 options: True, False
clamp_lenINT-1-1–2048
dropoutFLOAT0.10–0.9
batch_firstCOMBOTrue2 options: True, False
LAYER_STACKoptLIST

Outputs (1)

NameTypeDescription
LISTLIST