--- library_name: pytorch datasets: - bghira/minimax-music3-rvq-reverse-distillation tags: - audio - music - rvq - reverse-distillation - minimax-music-3 - mup --- # Open RVQ Encoder for MiniMax Music 3, 41M, v1 ## Status - V1 training is complete. - Recommended checkpoint: `checkpoint-17500`. - Not an official MiniMax model. - Not the original MiniMax Music 3 RVQ encoder. - No original encoder weights or source code were used. - Real-audio generalization is not established. - A packaged `from_pretrained` loader is not present yet. ## Objective Approximate the missing audio-to-RVQ path used by MiniMax Music 3. Input path: ```text 44.1 kHz waveform -> frozen DAV / Flow-VAE encoder -> 128-channel DAV latents -> this encoder -> 8 RVQ distributions per 25 Hz frame -> 1 semantic code + 7 acoustic codes ``` Output vocabularies: | Head | Role | Vocabulary | |---:|---|---:| | 0 | semantic | 16,384 | | 1-7 | acoustic | 1,024 each | The model predicts code distributions. Argmax produces a discrete code stream. The intended downstream test replays those codes through the MiniMax Music 3 LM, condition encoder, diffusion transformer, and DAV decoder. ## Architecture Exact trainable parameter count: **40,978,944**. | Component | Configuration | Parameters | |---|---|---:| | DAV latent input stem | Conv1d, 128 -> 512, kernel 7 | 459,264 | | Local residual stack | 3 blocks, dilations 1/3/9, GroupNorm, kernel-3 + kernel-1 convolutions | 3,151,872 | | Position embedding | learned, 128 x 512 | 65,536 | | Transformer | 8 pre-norm layers, width 512, 8 heads, FFN 2,048, GELU, dropout 0.1 | 25,219,072 | | Final normalization | LayerNorm(512) | 1,024 | | RVQ readouts | 8 independent `mup.MuReadout` heads | 12,082,176 | Processing order: 1. Apply the convolutional stem and residual stack at DAV latent rate. 2. Average-pool exact DAV latent spans into 25 Hz semantic frames. 3. Add learned positions. 4. Apply eight bidirectional Transformer encoder layers. 5. Apply final LayerNorm. 6. Produce one logit tensor per RVQ codebook. The pool matrix is supplied with each sample. It is not a fixed-ratio resampler. This preserves stitched-chunk alignment. Context: 128 semantic frames = 5.12 seconds. There is no cross-window state. ## Architecture Selection - The target size was set near 41M parameters. `Serveurperso` independently demonstrated that an encoder at this scale could preserve track and lyric identity through code replay. This implementation does not copy that encoder's weights or architecture. - DAV latents were selected instead of mel features. They are the continuous representation already used by the target pipeline. `marduk191`'s early mel proof of concept also showed the expected small-corpus generalization limit. - Convolutions handle local latent structure before temporal pooling. - The Transformer handles non-local interaction inside each 5.12-second crop. - Independent heads match the asymmetric semantic and acoustic vocabularies. - Width 512 and 8 heads give a fixed head dimension of 64. - Widths 128, 256, and 512 therefore map directly to 2, 4, and 8 heads. This is the μP width family. - Eight layers and FFN multiplier 4 place most capacity in temporal modeling while retaining a manageable DDP training cost. - A 128-frame context is the baseline, not a claimed optimum. A 256-frame follow-up is appropriate if semantic accuracy trails acoustic accuracy. ## Initialization and μP Package: [`microsoft/mup`](https://github.com/microsoft/mup). Shape family: | Model | Width | Heads | Head dimension | |---|---:|---:|---:| | base | 128 | 2 | 64 | | delta | 256 | 4 | 64 | | target | 512 | 8 | 64 | Initialization sequence: 1. Construct target, base, and delta models. 2. Call `mup.set_base_shapes(target, base, delta=delta)`. 3. Delete base and delta models. 4. Construct `mup.MuAdamW` after infshapes are attached. 5. Save `mup_base_shapes.bsh` with each exported checkpoint. Readouts: - All eight output layers are `mup.MuReadout`. - `output_mult = 1.0`. - `readout_zero_init = true`. - Readout weights and biases start at zero. - Initial output distributions are uniform within each vocabulary. Attention: - Score scale: `attention_multiplier / head_dim`. - `attention_multiplier = 8.0`. - Target scale: `8 / 64 = 1/8`, equal to standard `1/sqrt(64)` scaling. - Head dimension remains 64 across base, delta, and target widths. Other parameters: - Learned positions use `Normal(0, 0.02)`. - Convolution, attention, FFN, and normalization modules use their PyTorch initializers before μP shape metadata is attached. - Seed: 42, device-specific under DDP. μP supplies width-aware parameterization and optimizer scaling. The base/delta/target family supports μTransfer. The current `3e-4` learning rate is not presented as the result of a completed base-width hyperparameter sweep. ## Data Dataset: [`bghira/minimax-music3-rvq-reverse-distillation`](https://huggingface.co/datasets/bghira/minimax-music3-rvq-reverse-distillation). Run-launch snapshot: - 2,972 one-track ZIP shards. - 2,837 training records. - 135 holdout records. - Approximately 178 GB. - Synthetic tracks generated by MiniMax Music 3. - This is not MiniMax's original training set. Fields consumed by this trainer: - waveform audio; - sampled RVQ codes; - teacher top-50 token IDs; - teacher top-50 logits; - exact chunk-stitching metadata. The corpus also contains stored flow-VAE latents. This trainer does not consume them. It re-encodes waveform audio with [`SimpleTuner/MiniMax-Music-3-Encoder`](https://huggingface.co/SimpleTuner/MiniMax-Music-3-Encoder) and caches DAV latents once. Cached windows use `safetensors.safe_open(...).get_slice(...)`. Full-track latent tensors are not loaded for each crop. ## Alignment Nominal DAV ratio: `441 / 128 = 3.4453125` latents per semantic frame. The actual stitched timeline is not a global multiplication by that ratio. - Autoregressive rollout window: 200 semantic frames. - Rollout hop: 100 semantic frames. - Full stitched hop: 345 DAV latents. - Later chunks begin ownership 25 semantic frames after their nominal start. - Code row 0 is warm-up/priming. - Semantic frame `i` is supervised by code row `i + 1`. - The final partial chunk uses its own integer latent length. - Per-shard `chunk_stitching` bounds define the pool spans. - Training uses exact-alignment mode. Records without `chunk_stitching` metadata are excluded. These rules prevent cumulative label drift and training across incorrectly assigned rollout seams. ## Objective Function ```text loss = mean(CE_head_0 ... CE_head_7) + 0.25 * mean(KL_head_0 ... KL_head_7) ``` Hard targets: - Cross-entropy against sampled RVQ codes. - Equal weight for all eight heads. - Padding target: `-100`. Soft targets: - Teacher top-k: 50. - Temperature: 1.0. - Hinton `T^2` scaling. - Teacher distribution is renormalized over valid stored top-50 IDs. - Student uses full-vocabulary log-softmax, then gathers the teacher IDs. - Student probabilities are not renormalized over the top-50 subset. - Negative, EOS, and out-of-vocabulary teacher IDs are excluded. - Remaining teacher mass is renormalized after exclusion. - Frames with no valid teacher IDs are skipped for KL. The teacher logits come from LM predictions before audio-conditioned encoder output is available. Their uncertainty is useful but is not identical to an audio-conditioned posterior. This is why KL weight is 0.25 rather than 1.0. Equal head averaging is simple but imperfect. The semantic head has a much larger vocabulary and can dominate early CE. Per-head weighting is a possible follow-up. ## Training Run | Setting | Value | |---|---:| | Hardware | 4 x NVIDIA L40S | | Distribution | PyTorch DDP through Accelerate | | Precision | bfloat16 mixed precision | | Epochs | 20 | | Batch per rank | 16 | | Global batch | 64 | | Gradient accumulation | 1 | | Optimizer | `mup.MuAdamW` | | Learning rate | 3e-4 | | Weight decay | 0.01 | | LR schedule | SimpleTuner limited-data `cosine` | | Cosine half-cycle | 500 global steps | | Gradient norm limit | 1.0 | | Train crop | random 128-frame window | | Validation crop | deterministic 128-frame windows | | Validation interval | 500 steps | | Checkpoint interval | 500 steps | SimpleTuner's custom `cosine` schedule is intended for limited-data training. It repeatedly cools and reheats the learning rate, serving the same exploration purpose as SGDR warm restarts while moving smoothly through each cycle. A limited corpus can leave training stuck in a poor loss region when the learning rate selected for constant-with-warmup training is wrong. Repeated cycles give the optimizer additional chances to leave that region. V1 scheduler trace: - No linear warmup occurred. For this scheduler, `--lr_warmup_steps` sets the cosine interval. - The first μP parameter-group LR oscillated from `7.5e-5` to `1e-7` and back every 1,000 global steps. - Minima occurred at steps 500, 1,500, ..., 17,500. - Accelerate scheduler stepping maps the configured interval to this global-step cadence under four-rank DDP. Representative command: ```bash torchrun --standalone --nproc_per_node=4 scripts/train_minimax_music_rvq_encoder.py \ --dataset_repo_id bghira/minimax-music3-rvq-reverse-distillation \ --pretrained_vae_model_name_or_path SimpleTuner/MiniMax-Music-3-Encoder \ --latent_cache_dir cache/vae/minimaxmusic-rvq-encoder \ --output_dir output/minimaxmusic-rvq-encoder \ --require_exact_alignment \ --num_train_epochs 20 \ --train_batch_size 16 \ --mixed_precision bf16 \ --optimizer torch-adamw \ --learning_rate 3e-4 \ --weight_decay 0.01 \ --lr_scheduler cosine \ --lr_warmup_steps 500 \ --teacher_kl_weight 0.25 \ --teacher_kl_temperature 1.0 \ --window_frames 128 \ --window_stride 128 \ --d_model 512 \ --layers 8 \ --heads 8 \ --ff_mult 4 \ --dropout 0.1 \ --mup \ --mup_base_d_model 128 \ --mup_delta_d_model 256 \ --mup_readout_zero_init \ --checkpointing_steps 500 \ --validation_steps 500 \ --push_to_hub SimpleTuner/open-rvq-encoder-minimax-music3-41m-v1 ``` ## Checkpoint Format Each exported checkpoint contains: | File | Contents | |---|---| | `rvq_encoder.safetensors` | model state dictionary | | `rvq_encoder_config.json` | architecture and μP configuration | | `mup_base_shapes.bsh` | μP base-shape metadata | Trainer state, optimizer state, local paths, and credentials are not uploaded to this model repository. Loading currently requires the matching `RVQEncoderConfig` and `MiniMaxMusicRVQEncoder` definitions from `scripts/train_minimax_music_rvq_encoder.py`. ## Evaluation Current trainer metrics: - total validation loss; - hard CE; - teacher top-50 KL; - semantic top-1 token accuracy; - aggregate acoustic top-1 token accuracy. Four-rank real-data testing covered forward, backward, validation, checkpoint save, Hub export, and offline evaluation of every checkpoint. Full-run results are below. Machine-readable raw statistics: [`training-stats.json`](training-stats.json). It contains 882 training-log records and 36 full-holdout checkpoint evaluations. The JSON includes metric definitions and source-precision notes. Required end-to-end acceptance test: 1. Encode held-out waveform to DAV latents. 2. Predict eight codes per frame. 3. Replay predicted codes through the official LM path. 4. Compare replayed condition embeddings with stored condition embeddings. 5. Run the condition encoder, diffusion transformer, and DAV decoder. 6. Compare reconstructed audio and lyric identity with the source generation. Condition-embedding replay is implemented in a separate offline harness. It is not part of the trainer. Token top-1 is insufficient because multiple code sequences can be perceptually equivalent. Prior independent evidence from `Serveurperso`: - held-out STFT similarity: 0.83 to 0.87; - exact-code replay STFT similarity: 0.998; - acoustic exact-token match: 3% to 6%; - same music and lyrics remained identifiable after predicted-code replay; - a 550-track corpus overfit by epoch 17. Those numbers are from a separate encoder and training stack. They are not results for this checkpoint. ## Condition-Embedding Replay Comparison Protocol: - 130 exact-alignment holdout tracks; - each final checkpoint predicts argmax RVQ codes from cached DAV latents; - predicted codes are teacher-forced through the official language model and RVQ depth decoder; - hidden states pass through the official condition encoder with recorded chunk stitching; - reconstructed condition embeddings are compared with stored condition embeddings; - metric: per-track mean cosine over stitched condition-latent frames; - true sampled codes provide the replay control. | Model | Parameters | Mean cosine | Standard deviation | 5th-95th percentile | |---|---:|---:|---:|---:| | Serveurperso v1 | 40,978,944 | 0.663329 | 0.022175 | 0.628052-0.696328 | | SimpleTuner v1 | 40,978,944 | 0.762442 | 0.019550 | 0.734519-0.790450 | | SimpleTuner v2 | 154,736,064 | 0.769841 | 0.019063 | 0.742991-0.798636 | | SimpleTuner v3 | 154,736,064 | **0.770259** | 0.019274 | 0.741585-0.800492 | | True-code control | - | 0.999907 | - | - | ![Condition-embedding replay comparison](evaluation/replay-comparison/condition-embedding-replay-comparison.png) V1 exceeds the independent Serveurperso checkpoint by 0.099114 mean cosine. V2 adds 0.007399 over v1. V3 adds 0.000418 over v2. The MERT gain remains small downstream. This test stops before diffusion and DAV decode. It is not an STFT, waveform, lyric-identity, or listening score. Data: [`combined-aggregate.json`](evaluation/replay-comparison/combined-aggregate.json), [`provenance.json`](evaluation/replay-comparison/provenance.json), and [`raw per-record metrics`](evaluation/replay-comparison/raw-metrics-simpletuner-v1.json). ## Limitations - The final checkpoint regressed slightly from `checkpoint-17500` on the token-level holdout metrics. - 5.12-second encoder context. - No cross-window memory. - Synthetic model-output training domain. - Real audio is out of distribution until demonstrated otherwise. - Teacher uncertainty is from the LM rollout, not an audio-conditioned teacher encoder. - Exact token accuracy understates perceptual equivalence. - Semantic CE may dominate acoustic CE early. - Diffusion render and audio-domain evaluation remain pending. - Loading is not packaged as a stable library API. - Use is subject to the MiniMax Music 3 model terms and the reverse-distillation dataset terms. ## Discussion and Experimental Inputs Primary discussion: [MiniMaxAI/MiniMax-Music3 discussion #10, "Is the model trainable?"](https://huggingface.co/MiniMaxAI/MiniMax-Music3/discussions/10). Attribution below is for public discussion, measurements, datasets, and independent experiments. It does not imply shared authorship of this implementation. - [`bghira`](https://huggingface.co/bghira): ran the SimpleTuner training experiments; extracted sampled codes, teacher distributions, and alignment records; published the reverse-distillation corpus; organized this compatible-encoder run. - [`marduk191`](https://huggingface.co/marduk191): published WAV/code samples; built an early mel-based encoder proof of concept; reported small-corpus and real-audio limits; tested additional encoder variants. - [`scragnog`](https://huggingface.co/scragnog): calibrated HOT-Step CPP training against SimpleTuner; reported relative-weight-movement and loss measurements; identified structured-caption cache behavior and conditioning-rollout seam effects; confirmed SimpleTuner LoRA export interoperability with GGML. - [`Serveurperso`](https://huggingface.co/Serveurperso): independently built a 41M encoder, corpus generator, loader, and replay evaluation stack; demonstrated viable predicted-code replay; identified the stitched-hop, warm-up-row, and final-partial-chunk alignment rules. - [`dernet`](https://huggingface.co/dernet): explained why inference-time internal alignment does not provide target-derived alignment during training; clarified the role of RVQ token supervision; contributed tokenizer reverse-engineering analysis. Additional public artifacts: - [`marduk191/Minmax_music3_experiments` corpus](https://huggingface.co/marduk191/Minmax_music3_experiments/tree/main/corpus) - [`ServeurpersoCom/minimaxmusic.cpp` encoder proof commit](https://github.com/ServeurpersoCom/minimaxmusic.cpp/commit/d19efe9f94e41ac4c900aa30d56fe90c8dac7ef1) ## V1 Offline Checkpoint Evaluation Exact-alignment holdout: 130 tracks, 2,768 windows. | Selection | Checkpoint | Step | Loss | Semantic top-1 | Semantic top-5 | Acoustic top-1 | Acoustic top-5 | |---|---|---:|---:|---:|---:|---:|---:| | lowest loss; best semantic top-1; best acoustic top-1 | `checkpoint-17500` | 17,500 | 5.337856 | 0.4103 | 0.7838 | 0.0717 | 0.2094 | | final | `final` | 17,640 | 5.340557 | 0.4102 | 0.7835 | 0.0716 | 0.2092 | Top-k accuracy measures exact token inclusion. It does not measure perceptual code equivalence. ### Checkpoint Loss ![Checkpoint Loss](evaluation/v1/checkpoint-loss.png) ### Checkpoint Accuracy ![Checkpoint Accuracy](evaluation/v1/checkpoint-accuracy.png) ### Codebook Top1 ![Codebook Top1](evaluation/v1/codebook-top1.png) ### Training History ![Training History](evaluation/v1/training-history.png) Full data: [`checkpoint-metrics.csv`](evaluation/v1/checkpoint-metrics.csv), [`evaluation-metrics.json`](evaluation/v1/evaluation-metrics.json).