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

> Implement custom attention in Stable Diffusion's UNet by subclassing BasicTransformerBlock and replacing default layers to enhance image generation. Learn how to modify cross-attention for unique results.

- Repository: [CompVis - Computer Vision and Learning LMU Munich/stable-diffusion](https://github.com/CompVis/stable-diffusion)
- Tags: how-to-guide
- Published: 2026-03-02

---

**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`](https://github.com/CompVis/stable-diffusion/blob/main/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 `BasicTransformerBlock`s, and reshapes back to spatial tensors.

The `UNetModel` in [`ldm/modules/diffusionmodules/openaimodel.py`](https://github.com/CompVis/stable-diffusion/blob/main/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)`.

```python

# 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.

```python

# 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:

```python
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`](https://github.com/CompVis/stable-diffusion/blob/main/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`](https://github.com/CompVis/stable-diffusion/blob/main/ldm/modules/attention.py)**: Contains `CrossAttention`, `BasicTransformerBlock`, and `SpatialTransformer`—the core of the default attention pipeline.
- **[`ldm/modules/diffusionmodules/openaimodel.py`](https://github.com/CompVis/stable-diffusion/blob/main/ldm/modules/diffusionmodules/openaimodel.py)**: Implements `UNetModel`; the flag `use_spatial_transformer` determines whether cross-attention layers are inserted.
- **[`scripts/txt2img.py`](https://github.com/CompVis/stable-diffusion/blob/main/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.