Add RVQ encoder model card

README.md

@@ -0,0 +1,337 @@
+---
+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)
+