Stable Diffusion: Transformer2D's GroupNorm uses MLX's default eps (1e-5) where diffusers uses 1e-6

Summary

In stable_diffusion/unet.py, Transformer2D.__init__ builds its input GroupNorm without an eps:

self.norm = nn.GroupNorm(norm_num_groups, in_channels, pytorch_compatible=True)

mlx.nn.GroupNorm defaults to eps=1e-5. The reference implementation this mirrors, diffusers' Transformer2DModel, hardcodes 1e-6:

# diffusers/models/transformers/transformer_2d.py, _init_continuous_input
self.norm = torch.nn.GroupNorm(
    num_groups=self.config.norm_num_groups, num_channels=self.in_channels, eps=1e-6, affine=True
)

The epsilon is genuinely inconsistent across the UNet in diffusers, which is what makes it easy to miss: the resnets use 1e-5, the LayerNorms inside each transformer block use 1e-5, and only this spatial GroupNorm uses 1e-6. Taking the framework default unifies all three, which is tidier and does not match the reference.

This affects every model the example supports, since Transformer2D is shared: 16 modules in an SD 1.5-shaped UNet, and correspondingly more in SDXL.

Impact

Against reference activations dumped from the diffusers UNet on the same weights and the same inputs, the assembled UNet output moves by more than an order of magnitude:

  eps=1e-5 (current)   max_abs vs diffusers   5.628e-4
  eps=1e-6 (fixed)     max_abs vs diffusers   1.061e-5

The error is introduced at the first transformer and accumulates through the network. Comparing intermediates entry by entry, with the current default the first block's output is already 2.425e-4 off and the mid-block output reaches 6.208e-3; with eps=1e-6 those become 7.153e-6 and 1.224e-4. conv_in is unaffected either way, which is what localises it to the transformer rather than to weight loading or layout.

For calibration, an independent Rust implementation of the same UNet, verified against the same diffusers reference, sits at 1.1e-5 — so eps=1e-6 brings this example to the same agreement, and the current default is roughly 50x further out than a correct port.

It is a small absolute error and will not make images obviously wrong. It does put the example outside the tolerance anyone would use to verify a port against diffusers, and it costs nothing to fix: same operations, same shapes, one different scalar inside a normalisation that is already being computed.

Reproducer

No reference model or fixtures needed — run the UNet twice on identical random inputs, changing only the epsilon:

import sys
import mlx.core as mx
import numpy as np

sys.path.insert(0, ".")  # run from mlx-examples/stable_diffusion
from stable_diffusion.model_io import load_unet
from stable_diffusion.unet import Transformer2D

MODEL = "stabilityai/stable-diffusion-2-1-base"


def run(eps):
    unet = load_unet(key=MODEL, float16=False)
    n = 0
    if eps is not None:
        for _, m in unet.named_modules():
            if isinstance(m, Transformer2D):
                m.norm.eps = eps
                n += 1
    rng = np.random.default_rng(0)
    cross = unet.down_blocks[0].attentions[0].transformer_blocks[0].attn2.key_proj.weight.shape[1]
    x = mx.array(rng.standard_normal((1, 32, 32, 4)).astype(np.float32))
    t = mx.array(np.array([500.0], dtype=np.float32))
    ctx = mx.array(rng.standard_normal((1, 77, cross)).astype(np.float32))
    out = unet(x, t, encoder_x=ctx)
    mx.eval(out)
    return np.array(out, copy=False).astype(np.float64), n


default, _ = run(None)      # MLX default, eps=1e-5
fixed, n = run(1e-6)        # diffusers, eps=1e-6
d = np.abs(default - fixed)
print(f"Transformer2D modules affected: {n}")
print(f"output peak                   : {np.abs(fixed).max():.4f}")
print(f"max |eps=1e-5 - eps=1e-6|     : {d.max():.3e}")
print(f"mean|eps=1e-5 - eps=1e-6|     : {d.mean():.3e}")

Observed on an SD 1.5-shaped UNet at float32:

  Transformer2D modules affected: 16
  output peak                   : 3.9855
  max |eps=1e-5 - eps=1e-6|     : 6.249e-04
  mean|eps=1e-5 - eps=1e-6|     : 8.880e-05

(The measurements above were taken on SD 1.5, which is not in _MODELS and was added locally. Transformer2D is shared by every supported model, so stable-diffusion-2-1-base and sdxl-turbo are affected identically; the reproducer above is written against the default model.)

Fix

-        self.norm = nn.GroupNorm(norm_num_groups, in_channels, pytorch_compatible=True)
+        # diffusers' Transformer2DModel hardcodes eps=1e-6 here, while the
+        # resnets and the LayerNorms inside the blocks use 1e-5.
+        self.norm = nn.GroupNorm(
+            norm_num_groups, in_channels, pytorch_compatible=True, eps=1e-6
+        )

Happy to open a PR if useful.

Environment

• mlx 0.29.3, mlx-metal 0.29.3
• macOS (Darwin 25.6.0), arm64
• mlx-examples at current main