MIDI-RAE-JEPA-SON

TL;DR

We train a Swin-Transformer masked-autoencoder ("representation autoencoder") on piano-roll images of MIDI, using only internal-consistency objectives (LeJEPA attraction + SIGReg, factorization, cross-level masked-embedding prediction) — no reconstruction loss. This page collects the supplemental evidence that the learned six-level hierarchy encodes musical structure: probe suites over nine encoder variants, equivariance measurements, phrase-boundary detection against human annotations, and ablations of auxiliary harmonic supervision.

STORMBIRD probe results

STORMBIRD (Sweeping Test Of Representational Music Benchmarks, Information Retrieval & Diagnostics) is our suite of probes of musical qualities, available at github.com/drscotthawley/stormbird. Every encoder variant is evaluated with the full suite: linear probes on frozen embeddings at all six hierarchy levels (L0 = coarsest, one global vector per 128×128 crop; L5 = finest, 4×4-pixel patches). Cell shading is a per-row plasma heatmap (yellow = best in row); bold = best per row. Equivariance is reported as the R² of a linear fit of per-level embedding distance versus shift magnitude.

Encoder roster

Probe tables

Chord quality accuracy

Root note accuracy (12-class, chance .083)

Chroma regression R²

Key detection accuracy

Note density R²

Cross-song distance ratio (↓ lower is better)

Time-translation R²

EMOPIA 4-class accuracy (chance .25)

EMOPIA arousal (binary)

EMOPIA valence (binary)

Pitch transposition equivariance R²

Time translation equivariance R²

Structure analyses

Phrase-boundary detection vs. human annotations

We slide each encoder over full-song piano rolls, compute per-level self-similarity, extract Foote novelty curves, and score detected boundaries against the human-verified phrase annotations of Dai et al. (2020) (700 songs after alignment verification). F1 with tolerance in image columns (8 columns = 1 beat).

A linear probe asks whether a human-annotated boundary lies within half a bar of a 4-bar crop's center, from mean-pooled frozen embeddings per level (train/test split by song). Boundary information is real but weak, concentrated at coarse levels, and raw pixels carry none.

Linear boundary-detection probe — AP (chance = .17 (crop base rate))

Linear boundary-detection probe — AUC (chance = .50)

Equivariance curves

Per-level embedding distance as a function of pitch transposition / time translation, per encoder. Figures regenerate from the probe runs.

Pitch transposition

Time translation

Probe summaries

Soft factorization of pitch and time

A cosine-similarity loss on augmentation difference vectors pulls pitch- and time-shift directions toward parallel, anti-parallel, or orthogonal geometry depending on augmentation type; no directions are prescribed, only their pairwise relationships. The loss is applied only at the three coarsest levels (L0–L2).

Per level: PCA of normalized embedding difference vectors (left) and cosine-similarity histograms per pair type (right). The geometry visible at L3–L5 is inherited through the shared backbone rather than directly enforced.

Ablation studies

Design decisions behind the production recipe, from screening campaigns run during development: how transformer depth is allocated across the hierarchy, and how the main loss weights were chosen. Cells are shaded per row (plasma, yellow = best in row, honoring each row's arrow).

Ablations: model architecture

Eight alternative allocations of transformer blocks across the six Swin levels, trained and probed identically to the baseline (100-epoch screens, scalar λ=0.15, SIGReg on L0–L3). Content probes quoted at L5, their best level; ✗ = catastrophic L5 chroma collapse (R² of −77 to −713 despite normal training curves).

Depth buys little and risks a lot: the only decisive wins over the 16-block baseline are deep8's chroma and uniform3's root/key, while every variant with L2≥4 and L3=6 collapsed. The two best validation losses belong to a collapsed model and the worst-chroma survivor — val loss is a poor architecture-selection criterion here.

Ablations: hyperparameters

Per-level SIGReg strength (λ schedule)

Nine per-level λ schedules on the uniform3 backbone, all applying SIGReg to every level (n_skip=0); the reference applies scalar λ=0.15 to L0–L3 only. xmep1 adds cross-level masked-embedding prediction — the seed of the eventual production recipe.

No schedule beat the scalar reference on key/root — with the caveat, noted by the sweep itself, that the reference also skips SIGReg at L4/L5 entirely, so schedule shape and level count are confounded. The lightest schedule won val loss while its L5 cross-song ratio exceeded 1.0, and xmep1 won val loss while destroying L5 chroma: both dissociate loss from representation quality.

Chord-supervision weight λ_chord

Auxiliary chord heads (root + quality) at increasing loss weight, on coarse levels L0–L3 unless marked "/all" (all six levels; the per-level weight is the total divided by six). Values are each metric's best across levels, one consistent probe-suite vintage. The "/all" arms are 100-epoch screens; the rest are 250-epoch runs.

Chord supervision switches on between λ=0.1 (indistinguishable from baseline) and λ=0.5, with λ=1.0 marginally stronger. Spreading λ=1.0 over six levels cuts the effect to a third, and weight-matching (1.5/all = 0.25/level) restores it — the "all-levels hurts" reading was per-level dilution, not cross-level interference.

Listening examples

Each example takes a 4-bar crop from a held-out POP909 song (the Input tab), paints a blob-shaped mask over part of it, and asks the flow model to re-imagine the masked region four times (seeds 0–3) with conditioning dropout of 0.85 on the two finest hierarchy levels (L4/L5) — the demo's default “magic strength.” Music outside the mask is preserved exactly. Generated with the live demo's API (10 flow steps, guidance 1.0). Players need network access for the soundfont.

Runtime performance

Wall-clock execution time for the full pipeline (encode → sample → render) across hardware backends. Shading is a per-row plasma heatmap, inverted relative to the result tables: yellow = fastest. Parenthesized values are speedups relative to the 2-thread CPU baseline.

These are general-case timings (guidance strength ≠ 1.0, two function evaluations per Euler step). At the default strength of exactly 1.0 the unconditional branch is skipped and sampling times halve.

Links & citation

• POP909 dataset: music-x-lab/POP909-Dataset
• Phrase annotations: Dai, Zhang & Dannenberg (2020)
• EMOPIA: Hung et al. (2021)
• MIDI player: html-midi-player by Ondřej Cífka

@misc{midiraejepason2026,
  author = {Scott H. Hawley},
  title  = {Helping Music Co-Creation Agents `Listen' Well:
            Hierarchical Self-Supervised World Models for
            Understanding and Generation},
  note   = {Submitted to NeurIPS Creative AI Track},
  year   = {2026}
}