losslessmix-dit

DiT block merge (losslessmix for diffusion transformers)

Reduce the depth of a diffusion transformer (DiT) by merging pairs of adjacent residual blocks — without quantization and without training.

Idea

A DiT block is residual: x -> x + F(x). Two adjacent blocks compose as

B_j(B_i(x)) = x + F_i(x) + F_j(x + F_i(x))

Merging a pair into one block is lossy. The goal is to lose as little as possible:

Merge step — blend the two blocks' weights per channel by the absolute cosine similarity of their weight vectors:

k      = clamp(cos(W_i, W_j), 0, 1)
W      = W_i * k + W_j * (1 - k)

• Channels that point the same direction (high cos) blend proportionally;
• channels that point away (cos ≤ 0) keep W_j entirely.

So incompatible features are never averaged into garbage — that's the "lossless" part, inherited from recoilme/losslessmix.

Which pairs? Rank every block by block influence, measured on real activations:

BI_i = 1 - cos(x_i, x_{i+1})      # 0 = block is ~identity

Then greedily pick the emptiest (lowest combined BI) non-overlapping adjacent pairs, protecting the edge blocks (first/last and their neighbours), which are load-bearing.

Algorithm (3 steps)

1. analyze — run a few forward passes, hook each block's input/output, compute BI and pairwise S.
2. select — greedy by BI, non-overlapping, protect edges.
3. merge — merge_pair blends each pair into one block via cosine_merge_2d.

Results (Krea 2, 28-block single-stream DiT, bf16)

Beyond ~6 pairs (≈22 blocks) the degradation cliff appears — diffusion is iterative, so small per-block errors amplify over the denoising steps.

Usage

from diffusers import DiffusionPipeline
from merge_dit import analyze_blocks, select_pairs, apply_merges

pipe = DiffusionPipeline.from_pretrained("krea/Krea-2-Raw", torch_dtype=torch.bfloat16)
tr = pipe.transformer
blocks = tr.transformer_blocks

# 1. measure BI / S (run_forward runs one full transformer forward)
bi, S = analyze_blocks(tr, run_forward, n_passes=24)

# 2. pick pairs
pairs = select_pairs(bi, S, k=5, protect_head=2, protect_tail=2)

# 3. merge
apply_merges(blocks, pairs)

# update config and save / generate as usual
tr.register_to_config(num_layers=len(blocks))

Caveats

• Single-seed CLIP noise is ±0.04 — use several seeds for conclusions.
• The merge is a heuristic; it does not preserve the exact composition B_j(B_i(x)). A naive functional (least-squares) fit of the composition overfit on few samples and was worse than this weight merge.
• To go beyond ~17–20% of the DiT, you need fine-tuning / distillation.

Files

• merge_dit.py — the algorithm (merge + analyze + select + apply).