AudioSet MixIT Teacher: 475k/780k 50/50 Soup

This repository publishes a four-output, single-channel MixIT separator and the two trained checkpoints used to construct it. The default model is an inference-only 50/50 parameter interpolation of optimizer steps 475,000 and 780,000 from one AudioSet training trajectory. It was selected because it had the best aggregate result in a predeclared five-weight soup sweep on the full 514-example FUSS reverberant validation split.

This is a strong MixIT teacher release, not a claim that the complete SURF system or every number in the SURF paper has been reproduced. Compared with the paper's MixIT teacher row, this soup is lower on 1S and 2Si and slightly higher on 3Si and 4Si.

Results

All local results use the projected metrics from one fixed evaluation: 16 kHz, 160,000 samples, batch 2, seed 0, exact four-source selection, every one of the 514 examples in the FUSS reverberant validation split exactly once, and no W&B logging. The source-count suffix is the number of active reference sources in an example. 1S is SI-SDR; 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 four displayed source-count slices. The soup has zero non-finite outputs and a projected mixture-residual RMS of 5.29e-9. Under/equal/over source-count rates are 0.00389 / 0.34436 / 0.65175.

The exact soup, parent, and paper comparison values are preserved in evaluation/results_summary.json and the complete evaluator outputs are under evaluation/.

Use

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

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

Load the default soup 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 10-second model context, trims outputs to the original length, and writes source_1.wav through source_4.wav plus metadata.json. The four outputs are unordered latent sources. Audio longer than 160,000 samples is rejected because overlap, windowing, and permutation tracking are application decisions rather than properties of this checkpoint.

To load a parent instead of the soup:

from mixit_hf import MixITModel

separator = MixITModel.from_pretrained(
    checkpoint_filename="checkpoints/parents/step_0475000.pt",
    device="cuda",
)

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 MixIT 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 model repository. The AudioSet dataset card did not declare a license when this release was assembled, and the FUSS parquet card used the underspecified cc tag. Consult each upstream dataset and its source media terms before use. Exact local materialization identities, including every shard hash, are under training/data/, and immutable upstream revisions are in training/data/dataset_sources.json.

Training Recipe

The two parent checkpoints are from one run with these settings:

The learning rate 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 follows cosine decay from scale 0.3 to 0.05 for the final 200,000 updates. Step 475,000 was the validation leader; step 780,000 was a later point with stronger 2Si behavior. The run continued beyond both, but those later checkpoints did not dominate this soup's aggregate validation result.

The efficient hard solver computes the same discrete MixIT assignment objective while reducing assignment overhead; it is not a soft or approximate surrogate. Auxiliary diagnostics were logged, but their zero weights mean they did not affect gradients. Gradient accumulation was disabled because batch 18 per A100 fit the memory budget and avoided its throughput cost. A real preflight sweep tried per-rank batches 20, 18, 16, 14, 12, 10, and 8; batch 20 missed the registered free-VRAM margin, while batch 18 reached 105.292 examples/second.

The full machine-readable run contract and the human-readable recipe are in training/. The original production checkpoint records code_version=unknown because deployment did not include .git; the bundled inference modules are byte-for-byte copies of the archived source snapshot that loaded and evaluated these weights. The broader exact training source snapshot is not included in this model repository. The linked development repository's current public HEAD predates this checkpoint architecture, so it must not be treated as exact source provenance for resuming the run.

Soup Construction

The parents have identical model configuration, state-dict keys, tensor shapes, and dtypes. For each of 474 floating-point model tensors, the soup computes

theta_soup = 0.5 * theta_step475000 + 0.5 * theta_step780000

Interpolation was performed in float64 and cast back to each tensor's original dtype. Non-floating tensors, if any, are copied only after exact equality is verified. The five tested weights for the 475k parent were 0.1, 0.3, 0.5, 0.7, and 0.9; 0.5 produced the highest full-514 aggregate score.

The soup container inherits non-model metadata from the 780k checkpoint only so the existing evaluator can read it. Its optimizer, scheduler, RNG, step, and best-checkpoint fields do not describe a real training state. Therefore:

• use the soup for inference, teacher targets, or representation analysis;
• do not resume training from the soup;
• use one of the parent checkpoints when optimizer-state semantics matter.

The exact construction receipt, parent hashes, and sweep summary are in provenance/.

Why This Release

AudioSet supplies the large unlabeled mixture distribution needed by MixIT; FUSS supplies source-count-aware evaluation without leaking clean FUSS sources into training. Four outputs match the maximum FUSS source count and the SURF teacher contract. The 1M-step schedule was chosen because the shorter pilots had not reached the paper-level regime. Batch 18 per GPU maximized usable A100 throughput without gradient accumulation. The 50/50 soup was chosen from the measured sweep, not because averaging is assumed to help: it improved the 475k parent by 0.1937 dB aggregate and improved its 2Si, 3Si, and 4Si slices, while regressing 1S by 0.3227 dB.

The internal promotion gate still labeled this artifact continue: the aggregate gain was below its conservative +0.5 dB threshold and the 1S regression exceeded its -0.25 dB bound. Publication is for research use of the best observed aggregate MixIT teacher, not a claim that the gate or SURF SOTA criterion was cleared.

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 fixed receptive field is ten seconds. This release intentionally does not hide chunk boundaries or solve cross-window permutation tracking.
• Training audio inherits AudioSet/YouTube coverage and bias.
• The model is below the SURF paper MixIT row on 1S and especially 2Si.
• The parent checkpoints include optimizer and RNG state for provenance and research. The default soup is inference-only and non-resumable.
• 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.