Modulation¶

modulation ¶

Adaptive Layer Norm (AdaLN) modulation primitives for DiT-style models.

AdaLNModulation ¶

AdaLNModulation(dim: int, num_chunks: int, *, use_silu: bool = True, bias: bool = True)

Bases: Module

Project a conditioning vector to AdaLN modulation parameters.

Implements the standard SiLU → Linear projection used by DiT, PixArt, ERNIE-Image, LTX-Video, and friends. The output is a flat tensor of shape (B, num_chunks * dim) (or (B, S, num_chunks * dim) if conditioning is per-token); split it with mx.split(out, num_chunks, axis=-1) to recover the chunks.

Parameters:

Name	Type	Description	Default
`dim`	`int`	Model dimension. Both input and per-chunk output dim.	required
`num_chunks`	`int`	Number of modulation chunks the conditioning produces. Common values: `1` (single gate, e.g. AV cross-attention gate) `2` (shift + scale, e.g. text cross-attention) `4` (shift + scale, video + audio in AV models) `6` (MSA + MLP each shift / scale / gate — vanilla DiT) `9` (MSA + MLP + cross-attn — LTX-style)	required
`use_silu`	`bool`	Whether to pre-activate the input with SiLU before the Linear. Set `False` if the caller has already activated the conditioning (rare). Defaults to `True`.	`True`
`bias`	`bool`	Whether the Linear includes a bias. Defaults to `True`.	`True`

Weight keys

linear.weight, linear.bias (if bias=True).

ScaleShiftTable ¶

ScaleShiftTable(dim: int, num_params: int = 2)

Bases: Module

Final-layer scale/shift table.

The (num_params, dim) learnable table used by the final block-level AdaLN: rather than projecting an extra conditioning vector, the network stores a small parameter table that is added to a per-batch embedded timestep before splitting into scale/shift.

The forward signature is (B, dim) → (shift, scale). The two chunks are broadcast-ready for :func:modulate.

Parameters:

Name	Type	Description	Default
`dim`	`int`	Model dimension (per-chunk output dim).	required
`num_params`	`int`	How many chunks the table stores. Typically `2` (scale + shift) for the final pre-output norm.	`2`

Weight keys

table — (num_params, dim) learnable parameter.

gated_residual ¶

gated_residual(residual: array, gate: array, branch: array) -> array

Apply a gated residual update: residual + gate * branch.

modulate ¶

modulate(x: array, shift: array, scale: array) -> array

Apply AdaLN modulation: x * (1 + scale) + shift.

shift and scale broadcast over the sequence axis. For per-token modulation pass (B, S, dim); for shared modulation pass (B, dim) (or (B, 1, dim) for explicit broadcast).

adaln ¶

Adaptive Layer Norm (AdaLN) modulation primitives.

The DiT family (and its many spin-offs) modulates per-block features with shift/scale/gate parameters produced from a conditioning signal — typically the timestep embedding, sometimes pooled text. Each block consumes a tuple of chunks and applies them around its norm → attention/MLP → residual stages.

The pattern is small and uniform but every port re-implements it. This module exposes the building blocks:

:class:AdaLNModulation — SiLU → Linear → split into N chunks, the projection that turns a conditioning vector into modulation parameters. Supports any chunk count (1, 2, 4, 6, 9 are the common ones — single gate, scale+shift, AV cross-attn, MSA+MLP, MSA+MLP+CA).
:class:ScaleShiftTable — the (num_params, dim) learnable table used by the final layer of many DiTs. Indexed implicitly: it produces scale/shift for the post-block norm.
:func:modulate — x * (1 + scale) + shift with broadcasting.
:func:gated_residual — residual + gate * branch.

Caller composes these with mlx.nn.RMSNorm / LayerNorm and the attention / FFN of their choice. Time embedding lives in :mod:mlx_arsenal.diffusion; this module deliberately stays unaware of it, so any conditioning source (text pooler, class label, ...) can drive the modulation.

Example::

from mlx_arsenal.diffusion import TimestepEmbedding
from mlx_arsenal.modulation import (
    AdaLNModulation, modulate, gated_residual,
)

time = TimestepEmbedding(dim, dim)
modulation = AdaLNModulation(dim, num_chunks=6)
norm_msa = nn.RMSNorm(dim)

def block_forward(x, t_emb):
    params = modulation(time(t_emb))           # (B, 6*dim)
    shift_msa, scale_msa, gate_msa, *_ = params.split(6, axis=-1)
    h = modulate(norm_msa(x), shift_msa, scale_msa)
    h = attention(h)
    return gated_residual(x, gate_msa, h)

AdaLNModulation ¶

AdaLNModulation(dim: int, num_chunks: int, *, use_silu: bool = True, bias: bool = True)

Bases: Module

Project a conditioning vector to AdaLN modulation parameters.

Implements the standard SiLU → Linear projection used by DiT, PixArt, ERNIE-Image, LTX-Video, and friends. The output is a flat tensor of shape (B, num_chunks * dim) (or (B, S, num_chunks * dim) if conditioning is per-token); split it with mx.split(out, num_chunks, axis=-1) to recover the chunks.

Parameters:

Name	Type	Description	Default
`dim`	`int`	Model dimension. Both input and per-chunk output dim.	required
`num_chunks`	`int`	Number of modulation chunks the conditioning produces. Common values: `1` (single gate, e.g. AV cross-attention gate) `2` (shift + scale, e.g. text cross-attention) `4` (shift + scale, video + audio in AV models) `6` (MSA + MLP each shift / scale / gate — vanilla DiT) `9` (MSA + MLP + cross-attn — LTX-style)	required
`use_silu`	`bool`	Whether to pre-activate the input with SiLU before the Linear. Set `False` if the caller has already activated the conditioning (rare). Defaults to `True`.	`True`
`bias`	`bool`	Whether the Linear includes a bias. Defaults to `True`.	`True`

Weight keys

linear.weight, linear.bias (if bias=True).

ScaleShiftTable ¶

ScaleShiftTable(dim: int, num_params: int = 2)

Bases: Module

Final-layer scale/shift table.

The (num_params, dim) learnable table used by the final block-level AdaLN: rather than projecting an extra conditioning vector, the network stores a small parameter table that is added to a per-batch embedded timestep before splitting into scale/shift.

The forward signature is (B, dim) → (shift, scale). The two chunks are broadcast-ready for :func:modulate.

Parameters:

Name	Type	Description	Default
`dim`	`int`	Model dimension (per-chunk output dim).	required
`num_params`	`int`	How many chunks the table stores. Typically `2` (scale + shift) for the final pre-output norm.	`2`

Weight keys

table — (num_params, dim) learnable parameter.

modulate ¶

modulate(x: array, shift: array, scale: array) -> array

Apply AdaLN modulation: x * (1 + scale) + shift.

shift and scale broadcast over the sequence axis. For per-token modulation pass (B, S, dim); for shared modulation pass (B, dim) (or (B, 1, dim) for explicit broadcast).

gated_residual ¶

gated_residual(residual: array, gate: array, branch: array) -> array

Apply a gated residual update: residual + gate * branch.