AudioSet MixIT Teacher: Step 970k

This repository publishes a trained four-output, single-channel MixIT separator at optimizer step 970,000. It is an ordinary resumable training checkpoint, not a parameter soup. It was trained without labels on AudioSet and selected for downstream SURF research because it is the strongest trained checkpoint on the fixed FUSS validation protocol for both one-source SI-SDR and two-source SI-SDR improvement.

Step 925,000 remains the best aggregate and three-source checkpoint from the same trajectory. Step 970,000 is intentional: downstream use prioritizes accurately preserving a single source and resolving the common two-source case, while accepting modest regressions on three- and four-source examples.

This is a strong MixIT teacher release. It is not a claim that the complete SURF system or every SURF paper result has been reproduced.

Results

All trained-checkpoint results use one fixed evaluation protocol: 16 kHz, 160,000 samples, batch 2, seed 0, exact four-source output selection, the reverberant FUSS validation split, and all 514 examples exactly once. 1S is SI-SDR for examples with one active source. 2Si through 4Si are SI-SDR improvements over the input mixture.

The aggregate is the evaluator's dataset-level projected FUSS SI-SDR metric; it is not the arithmetic mean of the displayed source-count slices. Step 970k produced zero non-finite outputs, projected mixture-residual RMS 5.51e-9, and under/equal/over source-count rates of 0.00584 / 0.35214 / 0.64202.

The complete result, comparison, and checkpoint-pair receipts are under evaluation/ and provenance/.

Use

The release includes the exact inference-time architecture and a small Python SDK. Python 3.10 or newer is required.

python -m pip install \
  "mixit-audioset-fuss-step970k @ git+https://huggingface.co/sammyliu/mixit-audioset-fuss-step970k"

Load the default checkpoint directly from Hugging Face:

import torch

from mixit_hf import MixITModel

separator = MixITModel.from_pretrained(device="cuda")
mixture = torch.randn(2, 160_000)  # [batch, samples], mono 16 kHz
result = separator.separate_tensor(mixture)
sources = result.sources  # [batch, 4, samples]

Separate a mono or multichannel 16 kHz file no longer than ten seconds:

mixit-separate input.wav output_stems --device cuda

The file interface downmixes channels, pads short inputs to the ten-second context, trims outputs to the original length, and writes source_1.wav through source_4.wav plus metadata.json. Outputs are unordered latent sources. Audio longer than 160,000 samples is rejected because overlap, windowing, and permutation tracking are application-level decisions.

Only load PyTorch checkpoints from repositories you trust. These files use the standard Python-backed torch.load format.

Model

The checkpoint's model_config is authoritative. It is also extracted into training/checkpoint_embedded_config.json.

Training Data

Training used only unlabeled AudioSet audio from danjacobellis/audioset_opus_24kbps at revision a725d7cf1fea563c6eb9f6127dbd1d75b294668d. The materialized view contains the first 265,550 rows from 14 parquet shards, approximately 731 hours after decoding/resampling to 16 kHz and fixing examples at 160,000 samples. Each training example independently draws two AudioSet clips and sums them into a mixture-of-mixtures with the two original mixtures as references.

FUSS was not used for gradient updates. Validation used the reverberant validation split from scaleinvariant/fuss-parquet at revision b183027b802d1a488dea8b6f1d80daa32b76b681.

The audio corpora are not copied into this repository. Consult the upstream dataset cards and source-media terms before use. Exact materialization identities, including every shard hash, are under training/data/.

Training Recipe

The LR peaks at 1e-3. It warms linearly from zero for 5,000 updates, linearly anneals from scale 1.0 to 0.3 over the next 795,000 updates, then uses cosine decay from scale 0.3 to 0.05 for the final 200,000 updates. The run finished successfully at step 1,000,000; step 970,000 was selected after the trajectory completed.

The efficient hard solver computes the same discrete MixIT assignment objective while reducing assignment overhead; it is not a soft approximation. Auxiliary diagnostics were logged, but zero weights mean they did not affect gradients. Gradient accumulation was disabled because batch 18 per A100 fit and avoided its throughput cost. A measured preflight selected batch 18 at 105.292 examples/second after batch 20 missed the reserved-memory margin.

The full machine-readable run contract and human-readable recipe are in training/. The bundled inference modules are copied from the archived source snapshot that loaded and evaluated these weights. The broader distributed training source is not included, so the small SDK must not be treated as a complete resume implementation.

Selection Rationale

The entire late trajectory was evaluated every 5,000 updates when a durable checkpoint pair was available. Step 925k won aggregate and 3Si. Step 970k won both 1S and 2Si among trained checkpoints. The downstream SURF teacher was chosen explicitly for the latter profile:

• preserving one-source examples reduces pressure to invent extra sources;
• stronger 2Si targets the common two-source separation case;
• the measured 3Si/4Si tradeoff is accepted and remains visible in this card;
• no FUSS test examples were used to make the choice.

The preregistered conservative aggregate gate stopped at the one-million-step horizon because the terminal checkpoint did not clear its +0.5 dB aggregate threshold and slice-regression bounds. This post-training research selection does not rewrite that result. It records a different downstream utility choice.

Files And Integrity

Verify a downloaded checkout from its root:

shasum -a 256 -c SHA256SUMS

Limitations

• Outputs are anonymous and unordered; this is universal separation, not semantic stem extraction.
• The receptive field is ten seconds. Chunking and cross-window permutation tracking are intentionally left to the application.
• Training audio inherits AudioSet/YouTube coverage and bias.
• The model remains below the SURF paper MixIT row on 2Si and 4Si.
• Public download availability does not replace the terms in LICENSE.

References

• Wisdom et al., Unsupervised Sound Separation Using Mixture Invariant Training, NeurIPS 2020.
• Li et al., SURF: Separation via Unsupervised Remixing Flow, ICML 2026.
• Development implementation: samliu/mixture-invariant-training.