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

Open stable_diffusion/unet.py and inspect Transformer2D.init, where the input GroupNorm is constructed without an explicit epsilon. Run the provided reproducer from the stable_diffusion directory to compare the current output, then verify that the result matches the stated diffusers epsilon and improves agreement with the reference.