NNT DefineReformer Attention

Full self-attention has a nasty habit: every token looks at every token, so doubling the sequence length quadruples the compute. NNT DefineReformer Attention is the node that teaches you the workaround - the locality-sensitive hashing trick from the Reformer paper, where tokens are sorted into buckets and only attend within them. It's the pack's "efficient attention" lesson, and the layer it defines is a simplified take on that architecture.

What it actually does

It appends a ReformerAttention entry to the LAYER_STACK list that NntCompileModel compiles into a model. The core idea: instead of letting every position attend to every other position, you hash the query vectors so similar ones land in the same bucket, then restrict attention to tokens in the same bucket. The num_hashes setting runs this a few times with different hash functions so tokens that should be near each other don't get separated by bad luck. The result is attention that scales closer to linearly than quadratically with sequence length - that's the whole point of Reformer. The code comments in this pack are refreshingly honest here: the implementation is simplified and notes that real Reformer uses more machinery; this node is about grasping the concept.

Inputs that matter

• embed_dim - embedding width; 512 is the default.
• num_heads - attention heads, like any attention layer.
• num_buckets - how many hash buckets tokens get sorted into. Fewer buckets = cheaper but coarser grouping; 32 is the default.
• bucket_size - how many tokens can share a bucket. Together with num_buckets this bounds the per-bucket attention cost.
• num_hashes - the number of hash functions to run (redundancy so nothing important gets missed). 1 is the floor; 4–8 is typical.
• causal - True for autoregressive/language-model-style attention where a token only sees earlier positions.
• dropout, batch_first - the usual suspects; keep batch_first on.

The honest gotcha

This is the pack's most research-flavored corner, and the author's README frames the whole toolkit as a work in progress. The Reformer node produces a clean, inspectable definition that's excellent for teaching what LSH attention is, but the current compile path doesn't guarantee every transformer entry materializes in the compiled model yet - run NntCompileModel's "Only create script" mode to see what actually builds. If your actual goal is cheap long-context attention for a real project, this won't replace FlashAttention or a proper Reformer implementation. Its value here is conceptual: drag one in next to a vanilla attention node and compare the parameter math.

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. Heavy dependency list - torch, numpy, scikit-learn, pandas, transformers, shap - so expect a chunky first install. The pack's example workflows also want ComfyUI-Jjk-Nodes for text output; Manager's "Install Missing Custom Nodes" handles it.