minimax-music3-replanner-experiment

minimax-music3 latent re-planner: experiment log

This repo documents a line of experiments that did not produce a usable style-transfer model, and the measurements that explain why. We're publishing the checkpoints, training scripts, eval logs, and renders so nobody has to re-run this to find out what we found out. One sibling of this work did ship: the audio restoration variant became minimax-music3-latent-refiner-v0.10.

What we were trying to build

A small (137–143M parameter) flow-matching DiT that edits MiniMax Music 3 DAV latents. Input: a source track. Conditioning: MERT-v1-95M features of the source (all 13 hidden layers, injected per-layer n→n into the DiT blocks), a CLAP style vector (laion/larger_clap_music, trained on target audio embeddings so text embeddings can be swapped in at inference), optional RVQ codes, and a learned task embedding over {transfer, identity, restore}. Target: the DAV latents of a style-transferred version of the source. The pair corpus was 8,192 track pairs produced by an ACE-Step style-transfer worker (duration-matched source and restyled output), 64 pairs held out.

Model: single-stream DiT, d=768, 12 layers, 12 heads, RoPE over latent frames at 86.13 Hz (44100/512), flow matching with logit-normal timesteps. The in-context variants concatenate clean reference latents as extra tokens at shared RoPE positions with a learned segment embedding, compute loss only on the target half, and reattach the reference clean at every sampler step.

The arc

Each version below is a checkpoint in this repo. Numbers are holdout diagonal cosine (generated vs true target latents, N=32–64 confusion matrix) unless stated.

v1 (v1-baseline-perlayer) — scale ladder first: a 5.1M model overfits one pair to 0.9957 latent cosine in 2k steps, so the decoder path works. At 128 songs, per-layer MERT injection beats a same-capacity baseline by +0.16 diagonal at 4k steps. Full 8,192-pair DDP-8 run establishes the baseline.

v2 (v2-streams) — adds the CLAP style stream, RVQ code stream, MERT span masking with learned per-layer nulls. Stream ablation at ck7000: CLAP carries about half the holdout margin (0.0097 full vs 0.0046 without CLAP), the code stream changes nothing, MERT contributes little on the transfer task. That ablation was the first sign the transfer signal was thin.

v3/v4 (v4-curriculum) — correlation curriculum: mix identity pairs (input = target) and restore pairs (degraded input → clean target) with the transfer pairs, so the model sees tasks where input and output are actually correlated. v3 had a bug where the identity rate never reached the dataset constructor, so v4 is the first run where the curriculum was active.

The objective mistake. v1–v4 trained noise→target conditional generation: the model got MERT features of the source but never saw the source latents. SDEdit probes exposed it: pure generation on a training pair reached 0.121 cosine to the target — below the 0.135 you get by handing in the source latents unchanged. The model was learning to generate plausible music from MERT features, not to transport the source toward the target. Two corrected formulations were trained head-to-head from the same warm start:

v5b (v5b-bridge) — bridge transport: the flow runs from x1 = input latents (source, degraded, or self depending on task) to x0 = target latents. Task-loss signature was as expected — identity collapses instantly (zero-velocity anchor), restore fast, transfer slow. Holdout margin settled around +0.033–0.037 and stopped.

v6 (v6-incontext) — in-context editing: clean source latents concatenated as context tokens, generation from noise on the target half only. Climbed from +0.014 to +0.031 by 3k steps and kept the identity anchor clean. This became the main line.

v8 (v8-dpo-onpolicy-uncapped) — preference training, on-policy rejects: chosen = true target latents, rejected = the model's own 4-step generations for the same context, scored against a frozen copy of the warm-start weights, loss -log sigmoid(beta * margin) with beta=500. It destroyed itself within 1000 steps: the margin ran from 0 to 39 because raising your own error on your own generations is unbounded, and the rejected noised samples sit on the same distribution the sampler traverses at inference. Holdout diagonal fell to 0.052 while train losses stayed normal. The checkpoint is included because the failure mode is invisible in the training loss and worth having as a reference artifact.

v9 (v9-dpo-onpolicy-capped) — same recipe with the positive margin capped at 1.0 and beta=10. Contained (diagonal 0.115 at 1k instead of 0.05) but still below the no-preference baseline. Bounded or not, spending gradient on raising error over your own sampling trajectory costs sampling quality.

v10 (v10-dpo-source-reject) — rejected = the source latents themselves, i.e. the preference pair is "produce the target" vs "copy the input". No sampling needed, no self-referential loop, and the pressure aims at the copy attractor specifically. Best of the three DPO variants: margin +0.023 at 1k with the sampler intact, holdout residual cosine 0.40–0.46. Then 13,000 further steps changed nothing. Eval trajectory is in evals/eval_v10.jsonl: diagonal margin oscillates between +0.011 and +0.023 from step 1000 to step 14000, residual cosine oscillates 0.34–0.46, no trend. Stopped at 14,450 steps, 112 epochs.

Why it plateaus

The pair corpus doesn't contain much learnable transfer signal, and we measured this directly:

• source↔target latent cosine across the ACE pairs: 0.135. The style transform is close to arbitrary per pair at the latent level.
• frame-wise MERT correlation between source and target: 0.233, against 0.16 for shuffled pairs. About +0.07 of alignable structure for a frame-conditioned model to grab.

Every objective change after that measurement (bridge, in-context, three DPO variants) moved the needle inside a band of a few hundredths and none escaped it. The restoration task, where input and output are strongly correlated by construction, worked fine with the same architecture and became the refiner release. The conclusion we're acting on: the next attempt needs pairs with high mutual information — stem swaps, or MIDI transcribe-and-rerender pairs where the target is genuinely predictable from the source — not another objective.

Audio examples

30-second crops, in-context sampler, 32 Euler steps, v10 checkpoint. Three inputs: a pair from the training split, a pair from the 64-pair holdout, and Take Five (Dave Brubeck, 1959) — instrumental jazz, nothing remotely like it in the corpus, included to show behaviour with no coverage at all.

The Take Five reference crops are excerpts of a commercial recording, so only the model outputs are included for that input.

Expected reading: identity holds up everywhere (that anchor trained to ~0.005 loss), restore is usable on in-domain material, transfer is loosely related to its target even on the training split, and Take Five comes back as something else entirely — 44.1k pop-adjacent texture is all the model has ever produced.

Contents

• checkpoints/<version>/model.safetensors — bf16 model weights (no optimizer state), plus the training args as recorded in the original state.pt.
• scripts/ — the exact training and data scripts: train_latent_replanner.py (model, sampler), train_latent_replanner_ddp.py (DDP trainer: pair dataset, online MERT/CLAP/DAV extraction, degradation chain, all objectives incl. the DPO variants), precompute_rvq_codes.py, render_taskspace.py (produced the examples above).
• evals/eval_v{8,9,10}.jsonl — per-1k-step holdout metrics as logged during training. Logs for v1–v6 did not survive their shells; the numbers quoted above were recorded at the time.
• normalization_stats.pt — latent and MERT-layer mean/std used for training and sampling.

Loading a checkpoint

import torch
from safetensors.torch import load_file
from train_latent_replanner import LatentReplanner, sample

model = LatentReplanner(128, 768, 768, 12, 12)
model.enable_layer_conditioning(768, 13)
model.enable_mert_masking(768, 13)
model.enable_style_conditioning(512)
model.enable_context_editing()
model.enable_task_conditioning(3)
model.enable_code_conditioning(23552, 8)
model.load_state_dict(load_file("checkpoints/v10-dpo-source-reject/model.safetensors"))

Earlier checkpoints enable fewer streams; each training_args.json records which flags were active. Conditioning extraction and normalization follow render_taskspace.py.