How to Implement Custom Attention Mechanisms in Stable Diffusion's UNet

You can replace the default cross-attention layers in Stable Diffusion by subclassing BasicTransformerBlock to inject your custom attention class, then replacing the transformer blocks inside the UNet's SpatialTransformer modules while preserving the original forward API.

Stable Diffusion's UNet relies on cross-attention mechanisms to condition the diffusion process on text embeddings. If you need to experiment with linear attention, Performer, or other efficient variants, you'll need to modify the architecture in CompVis/stable-diffusion while keeping the pipeline compatible with existing training and inference scripts.

Understanding the Default Cross-Attention Architecture

The conditioning capability flows through three tightly coupled classes defined in ldm/modules/attention.py:

• CrossAttention (line 52): The standard multi-head attention implementation that operates as self-attention when context is None and cross-attention when text embeddings are provided.
• BasicTransformerBlock (line 96): Wraps two CrossAttention instances (attn1 for self-attention, attn2 for cross-attention) plus a feed-forward network.
• SpatialTransformer (line 118): Converts 2D feature maps to sequences, processes them through one or more BasicTransformerBlocks, and reshapes back to spatial tensors.

The UNetModel in ldm/modules/diffusionmodules/openaimodel.py instantiates these components when use_spatial_transformer=True, inserting SpatialTransformer blocks at resolutions specified by attention_resolutions. During the forward pass, the context tensor (typically CLIP text embeddings) flows into attn2, enabling text-to-image conditioning.

Because BasicTransformerBlock constructs CrossAttention directly in its __init__ method (lines 99-103), you cannot configure the attention class via constructor arguments. Instead, you must override the instantiation logic through subclassing or monkey-patching.

Strategies for Replacing Attention Mechanisms

You have two primary approaches to inject custom attention:

1. Subclassing: Create a descendant of BasicTransformerBlock that constructs your custom attention module instead of the default CrossAttention, then recursively replace blocks inside the UNet. This is maintainable and explicit.

2. Monkey-patching: Override the CrossAttention class reference at import time before the UNet is constructed. This is faster for experiments but harder to maintain in production code.

Both strategies preserve the rest of the diffusion pipeline because they maintain the exact signature forward(x, context=None, mask=None) that the UNet expects.

Creating a Custom Attention Module

Define a module that mirrors the interface of the original CrossAttention. The class must accept query_dim, context_dim, heads, dim_head, and dropout in its constructor, and implement forward(x, context=None, mask=None).

# custom_attention.py

import torch
import torch.nn as nn
from einops import rearrange
from ldm.modules.attention import default, exists, max_neg_value

class MyCrossAttention(nn.Module):
    """
    Example: a lightweight linear-attention variant that keeps the same
    API as the original CrossAttention (q, k, v → out).
    """
    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0):
        super().__init__()
        inner_dim = dim_head * heads
        context_dim = default(context_dim, query_dim)

        self.heads = heads
        self.scale = dim_head ** -0.5

        self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
        self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
        self.to_v = nn.Linear(context_dim, inner_dim, bias=False)

        self.to_out = nn.Sequential(
            nn.Linear(inner_dim, query_dim),
            nn.Dropout(dropout)
        )

    def forward(self, x, context=None, mask=None):
        h = self.heads
        q = self.to_q(x)                       # (b, n, h·d)

        context = default(context, x)
        k = self.to_k(context)
        v = self.to_v(context)

        # reshape to (b·h, n, d)

        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=h), (q, k, v))

        # **Linear attention**: softmax on keys only, then a weighted sum of values

        k = torch.softmax(k, dim=1)            # (b·h, n, d)

        attn = torch.einsum('b i d, b j d -> b i j', q, k) * self.scale
        out = torch.einsum('b i j, b j d -> b i d', attn, v)

        out = rearrange(out, '(b h) n d -> b n (h d)', h=h)
        return self.to_out(out)

Integrating Custom Attention into the UNet

Subclass BasicTransformerBlock to override the attention constructors, then create a UNet wrapper that swaps out the original transformer blocks for your custom implementations.

# custom_unet.py

import torch.nn as nn
from ldm.modules.diffusionmodules.openaimodel import UNetModel
from ldm.modules.attention import BasicTransformerBlock
from custom_attention import MyCrossAttention

class MyTransformerBlock(BasicTransformerBlock):
    """Override only the attention constructors."""
    def __init__(self, dim, n_heads, d_head, dropout=0., context_dim=None, gated_ff=True, checkpoint=True):
        super(nn.Module, self).__init__()
        # Replace the two CrossAttention instances with MyCrossAttention

        self.attn1 = MyCrossAttention(query_dim=dim, heads=n_heads, dim_head=d_head, dropout=dropout)      # self-attention

        self.ff    = self.ff  # keep the default FFN

        self.attn2 = MyCrossAttention(query_dim=dim, context_dim=context_dim,
                                      heads=n_heads, dim_head=d_head, dropout=dropout)               # cross-attention

        self.norm1 = nn.LayerNorm(dim)
        self.norm2 = nn.LayerNorm(dim)
        self.norm3 = nn.LayerNorm(dim)
        self.checkpoint = checkpoint

# A thin wrapper around UNetModel that injects MyTransformerBlock

class UNetWithMyAttention(UNetModel):
    def __init__(self, *args, **kwargs):
        # Enable the spatial transformer so we can replace its blocks

        kwargs.setdefault('use_spatial_transformer', True)
        super().__init__(*args, **kwargs)

        # Replace every BasicTransformerBlock inside the SpatialTransformer

        for block in self.input_blocks:
            self._replace_transformer(block)

        for block in self.middle_block:
            self._replace_transformer(block)

        for block in self.output_blocks:
            self._replace_transformer(block)

    def _replace_transformer(self, module):
        if isinstance(module, nn.ModuleList):
            for sub in module:
                self._replace_transformer(sub)
        elif hasattr(module, 'transformer_blocks'):
            # module is a SpatialTransformer

            new_blocks = nn.ModuleList([
                MyTransformerBlock(
                    dim=tb.attn1.to_q.in_features,
                    n_heads=self.num_heads,
                    d_head=tb.attn1.to_q.out_features // self.num_heads,
                    dropout=0,
                    context_dim=self.context_dim,
                    gated_ff=True,
                    checkpoint=self.use_checkpoint,
                )
                for tb in module.transformer_blocks
            ])
            module.transformer_blocks = new_blocks

You can now instantiate the modified UNet without changing any downstream training or inference code:

from custom_unet import UNetWithMyAttention

model = UNetWithMyAttention(
    image_size=64,
    in_channels=4,
    model_channels=320,
    out_channels=4,
    num_res_blocks=2,
    attention_resolutions=(4, 2, 1),
    dropout=0.0,
    channel_mult=(1, 2, 4, 8),
    use_spatial_transformer=True,    # required

    context_dim=768,                # e.g. CLIP text embeddings

    num_heads=8,
    num_head_channels=-1,
)

The model remains compatible with scripts/txt2img.py and other pipeline tools because the public API (forward(x, timesteps, context)) is unchanged.

Key Source Files for Reference

• ldm/modules/attention.py: Contains CrossAttention, BasicTransformerBlock, and SpatialTransformer—the core of the default attention pipeline.
• ldm/modules/diffusionmodules/openaimodel.py: Implements UNetModel; the flag use_spatial_transformer determines whether cross-attention layers are inserted.
• scripts/txt2img.py: Demonstrates how the UNet is instantiated and how context tensors (CLIP embeddings) are passed during inference.

Summary

• Cross-attention in Stable Diffusion is implemented via CrossAttention classes inside BasicTransformerBlock modules, which are stacked within SpatialTransformer layers in the UNet.
• To implement custom attention mechanisms, subclass BasicTransformerBlock to instantiate your custom class instead of the default CrossAttention.
• Replace the transformer blocks recursively in the UNet's input_blocks, middle_block, and output_blocks while preserving the original method signatures.
• This approach maintains full compatibility with existing Stable Diffusion pipelines, allowing you to experiment with linear attention, Performer, or other variants without breaking training scripts.

Frequently Asked Questions

Can I use monkey-patching instead of subclassing to replace cross-attention?

Yes. You can reassign the CrossAttention class reference in ldm.modules.attention before the UNet is instantiated, causing BasicTransformerBlock to construct your custom class automatically. This works for rapid prototyping but is harder to maintain than explicit subclassing because it relies on import order and global state.

What parameters must my custom attention module accept?

Your constructor must accept query_dim, context_dim (optional), heads, dim_head, and dropout. The forward method must accept x (the feature tensor), context (optional conditioning), and mask (optional attention mask), returning a tensor of the same shape as x. Matching this interface ensures compatibility with BasicTransformerBlock.

Will custom attention break existing training checkpoints?

If your custom attention module changes the state dictionary keys (for example, by renaming to_q, to_k, to_v projections), you will not be able to load weights from the original Stable Diffusion checkpoints into your custom layers. To resume training from pretrained weights, either map the old keys to your new architecture during loading or keep the parameter names identical to the original CrossAttention implementation.

Which UNet resolutions use spatial transformers?

The attention_resolutions parameter in UNetModel.__init__ (typically set to (4, 2, 1)) controls which downsampling levels receive SpatialTransformer blocks. A value of 4 applies attention at the highest resolution (least downsampled), while 1 applies it at the lowest resolution. You can verify which blocks contain transformers by checking for the transformer_blocks attribute in your model's modules.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}