# Architecture of the X-Transformers Encoder in Stable Diffusion: Core Components and Data Flow

> Explore the x-transformers encoder architecture in Stable Diffusion. Understand its core components like self-attention and feed-forward layers and their data flow within the model.

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

---

**The x-transformers encoder in Stable Diffusion is a bidirectional transformer implemented in [`ldm/modules/x_transformer.py`](https://github.com/CompVis/stable-diffusion/blob/main/ldm/modules/x_transformer.py) that stacks self-attention and feed-forward layers within an `AttentionLayers` backbone, wrapped by `TransformerWrapper` to handle embeddings and positional encoding.**

The CompVis/stable-diffusion repository leverages the implementation from lucidrains/x-transformers to power its text-conditioning pipeline. This encoder serves as the backbone for processing tokenized text prompts into latent embeddings that guide the diffusion process. Understanding its modular architecture reveals how the model handles sequence data through configurable attention mechanisms, normalization schemes, and residual connections.

## Core Building Blocks

The encoder is constructed from several interconnected classes defined in [`ldm/modules/x_transformer.py`](https://github.com/CompVis/stable-diffusion/blob/main/ldm/modules/x_transformer.py). Each component handles specific aspects of the transformation pipeline.

### Attention Mechanism

The **`Attention`** class implements multi-head scaled dot-product attention with extensive customization options. It supports **talking heads** (learned linear projections before and after softmax), sparse top-k attention, and memory key/value vectors via `num_mem_kv`. The class also handles causal masking for autoregressive tasks and integrates rotary or relative positional embeddings. In the context of the encoder used for text conditioning, causality is disabled to allow bidirectional context flow.

### FeedForward Layers

The **`FeedForward`** class provides a two-layer MLP with optional gating mechanisms. By default, it uses a standard GELU activation, but can be configured to use **gated-GEGLU** variants for improved expressiveness. Dropout is applied for regularization between the linear transformations.

### AttentionLayers Stack

**`AttentionLayers`** serves as the foundational transformer stack that sequences attention and feed-forward operations. It accepts a `layer_types` schedule that determines whether each layer performs self-attention (`'a'`), cross-attention (`'c'`), or feed-forward processing (`'f'`). By default, the encoder alternates between attention and feed-forward blocks (`('a', 'f')`). This class supports advanced architectural patterns including pre-norm/post-norm configurations, ScaleNorm or RMSNorm replacements, ReZero residuals, GRU gating, macaron structures, and Parallel Attention Reparameterisation (PAR).

### Encoder Class

The **`Encoder`** class is a thin subclass of `AttentionLayers` configured specifically for bidirectional processing. It sets `causal=False` and restricts operations to self-attention only, making it suitable for understanding complete sequence context rather than generating sequences autoregressively. The class definition appears at lines 41-45 in [`ldm/modules/x_transformer.py`](https://github.com/CompVis/stable-diffusion/blob/main/ldm/modules/x_transformer.py).

### TransformerWrapper

**`TransformerWrapper`** acts as the high-level interface that prepares inputs for the encoder. It manages token embeddings (`nn.Embedding`), optional absolute positional embeddings, learned memory tokens (similar to CLS tokens), and dimension projection layers. After processing through the encoder stack, it applies final normalization and optionally projects to vocabulary logits.

## Data Flow Through the Encoder

The transformation pipeline follows a strict six-stage process when processing text tokens:

1. **Token Embedding**: Input token IDs pass through `self.token_emb` to create dense vector representations.
2. **Positional Encoding**: If `use_pos_emb` is enabled and the encoder lacks internal positional support, `AbsolutePositionalEmbedding` adds position information.
3. **Memory Token Injection**: Optional learned memory tokens prepend to the sequence, similar to BERT's CLS token.
4. **Dimension Projection**: `self.project_emb` aligns embedding dimensions with the encoder's internal `dim`.
5. **Attention Processing**: The sequence flows through `AttentionLayers`, where each block applies (pre-norm) normalization followed by either self-attention or feed-forward processing, then residual connections.
6. **Output Normalization**: Final layer normalization (`self.norm`) produces the hidden states used for conditioning.

## Configurability Options

The encoder exposes numerous architectural flags through the `AttentionLayers` constructor:

- **`use_scalenorm`** or **`use_rmsnorm`**: Replaces standard LayerNorm with ScaleNorm or RMSNorm variants.
- **`use_rezero`**: Implements ReZero residual connections (`g * f(x) + x`) where `g` is a learned scalar.
- **`gate_residual`**: Substitutes additive residuals with GRU-gated residual connections.
- **`macaron`**: Inserts additional feed-forward layers before attention blocks (the "macaron" architecture).
- **`sandwich_coef`**: Adds extra attention layers at the beginning and end of the stack.
- **`par_ratio`**: Enables Parallel Attention Reparameterisation scheduling.
- **`talking_heads`**: Activates learned linear projections on attention heads before and after softmax.
- **`sparse_topk`**: Restricts attention to the top-k highest affinity keys.
- **`num_mem_kv`**: Adds learned memory key/value vectors concatenated to attention inputs.

## Usage in Stable Diffusion's Text Conditioning

Stable Diffusion instantiates the encoder through several wrapper classes in [`ldm/modules/encoders/modules.py`](https://github.com/CompVis/stable-diffusion/blob/main/ldm/modules/encoders/modules.py):

- **`TransformerEmbedder`**: Direct token-to-embedding pipeline using the x-transformers encoder.
- **`BERTEmbedder`**: Combines BERT tokenization with the x-transformers architecture.
- **`FrozenCLIPEmbedder`**: Uses CLIP's native transformer (not the x-transformer implementation).

These wrappers create a `TransformerWrapper` instance with `attn_layers=Encoder(dim=n_embed, depth=n_layer)`, connecting the x-transformers backbone to the U-Net's cross-attention mechanisms via the `n_embed` and `n_layer` parameters specified in configuration files like [`configs/stable-diffusion/v1-inference.yaml`](https://github.com/CompVis/stable-diffusion/blob/main/configs/stable-diffusion/v1-inference.yaml).

## Code Examples

### Basic Encoder Usage via TransformerEmbedder

```python
import torch
from ldm.modules.encoders.modules import TransformerEmbedder

# Initialize embedder with x-transformers encoder backbone

embedder = TransformerEmbedder(
    n_embed=768,
    n_layer=12,
    vocab_size=49408,  # OpenAI CLIP vocabulary size

    max_seq_len=77,
    device="cpu",
)

# Process dummy token IDs

tokens = torch.randint(0, 49408, (1, 77))
embeddings = embedder(tokens)  # Shape: [1, 77, 768]

print(embeddings.shape)  # torch.Size([1, 77, 768])

```

### BERT-Style Encoder with Tokenization

```python
from ldm.modules.encoders.modules import BERTEmbedder

bert_encoder = BERTEmbedder(
    n_embed=512,
    n_layer=8,
    vocab_size=30522,  # BERT base vocabulary

    max_seq_len=77,
    device="cpu",
    use_tokenizer=True,
    embedding_dropout=0.1,
)

# Automatic tokenization and encoding

text = ["a photo of a sunny beach"]
z = bert_encoder(text)  # Shape: [1, 77, 512]

print(z.shape)

```

### Extracting Attention Maps for Visualization

```python

# Access attention weights for debugging

logits, attn_maps = embedder.transformer(
    tokens,
    return_attn=True,
)

# attn_maps is a list of tensors, one per layer

# Each tensor shape: [batch, heads, seq_len, seq_len]

print(len(attn_maps), attn_maps[0].shape)

```

## Summary

- The x-transformers encoder in Stable Diffusion resides in [`ldm/modules/x_transformer.py`](https://github.com/CompVis/stable-diffusion/blob/main/ldm/modules/x_transformer.py) and inherits from the lucidrains/x-transformers library.
- Core components include the `Attention` class for multi-head attention, `FeedForward` for MLP processing, and `AttentionLayers` for stack management.
- The `Encoder` subclass configures `AttentionLayers` for bidirectional (non-causal) processing with only self-attention.
- `TransformerWrapper` handles embeddings, positional encoding, and memory tokens before passing data to the encoder.
- Extensive configurability options include ScaleNorm/RMSNorm, ReZero residuals, GRU gating, macaron architecture, and talking heads.
- Stable Diffusion uses this encoder through `TransformerEmbedder` and `BERTEmbedder` wrappers to convert text prompts into conditioning vectors.

## Frequently Asked Questions

### What is the difference between the Encoder and AttentionLayers classes?

The `AttentionLayers` class is the generic transformer stack that handles layer scheduling and residual connections, supporting both causal and non-causal modes. The `Encoder` class is a specific subclass that hardcodes `causal=False` and configures the layer schedule for pure self-attention without autoregressive masking, making it suitable for understanding complete text sequences.

### Where does the x-transformers encoder get its positional information?

Positional encoding occurs at the `TransformerWrapper` level through `AbsolutePositionalEmbedding` added to token embeddings, unless the underlying `AttentionLayers` is configured with `position_infused_attn` or rotary embeddings. The wrapper checks `use_pos_emb` and the encoder's internal capabilities before applying external positional encoding.

### Can I use sparse attention with the Stable Diffusion text encoder?

Yes, the `Attention` class supports `sparse_topk` which restricts attention to the top-k highest affinity keys. However, the default Stable Diffusion configuration typically uses full attention. To enable sparse attention, you would need to instantiate the encoder with `sparse_topk` set in the `AttentionLayers` constructor parameters passed to `TransformerEmbedder`.

### How does the encoder handle variable-length sequences?

The encoder processes fixed-length sequences padded to `max_seq_len` (typically 77 tokens). The `TransformerWrapper` manages padding through its embedding layer, and while the attention mechanism computes full matrix multiplications, the effective receptive field is limited to the actual token positions. No explicit masking for variable lengths is applied in the default implementation beyond the fixed sequence dimension.