---
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 WIP

## Status

- Work in progress.
- Checkpoints are uploaded during training.
- 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 | cosine |
| Warmup | 500 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 |

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-music-3-41m-v1-wip
```

## 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 smoke testing covered forward, backward, validation, checkpoint save, and Hub export. Full-run results will be added after checkpoints are evaluated.

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.

This condition-embedding/replay evaluation is not implemented in the current 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.

## Limitations

- WIP weights can regress between checkpoints.
- 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.
- End-to-end condition-embedding evaluation is 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)