Skip to content

Add opt-in fresh DiT adaLN-Zero initialization - #9

Open
JiamingPan wants to merge 1 commit into
mainfrom
codex/dit-adaln-zero-init
Open

JiamingPan wants to merge 1 commit into
mainfrom
codex/dit-adaln-zero-init

Conversation

@JiamingPan

Copy link
Copy Markdown
Collaborator

Adds an opt-in model.initialization: adaln_zero policy for freshly constructed native DiTs.

  • Zero block conditioning projections (norm1.linear), final conditioning (proj_out_1), and prediction head (proj_out_2).
  • Leave attention, FFN, embeddings, architecture, loss, and samplers unchanged.
  • Preserve native initialization by default; never reapply initialization when loading or resuming checkpoints.
  • Validate the full projection API before changing weights and reject unsupported models or repeated initialization.
  • Document usage and add 19 regression tests.

Validation: all 97 repository tests pass locally on CPU with diffusers 0.38.0, PyTorch 2.5.1, and Python 3.12.7, including identity blocks, trainable output heads, and checkpoint preservation.

This ports the zero-modulation/output policy tested in the L16 ablation, not the complete reference DiT recipe or a guarantee of generalization.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant