Source-linked AI summary
Helping Music Co-Creation Agents 'Listen' Well: Hierarchical Self-Supervised World Models for Understanding and Generation
Scott H. Hawley
TL;DR
Collaborative music agents need representations that support both musical understanding and generation while preserving human agency. This paper develops a self-supervised hierarchical symbolic-music model and finds that representation levels align with musical time scales, while generation supports controllable suggestions and inpainting.
Problem
Collaborative music agents need interpretable musical representations for feedback and suggestions, but labeled hierarchical music data is limited and reconstruction and representation learning have been treated as separate tasks.
Method
The system learns hierarchical symbolic-music representations without labels or music-theory vocabulary, then uses pixel-space flow matching conditioned on compressed embeddings for controllable generation and inpainting.
Results
Musical properties become decodable at levels matching their time scales: phrase boundaries at coarse levels, and note density and harmonic content at fine levels.
Takeaways & Limitations
The system supports a collaborative workflow in which the agent offers verbal feedback and limited musical suggestions while leaving implementation to the human.
Takeaways & Limitations
Inpainted notes generally fit surrounding music but follow it less closely than larger diffusion models, with some notes landing off.
Abstract
from arXiv · showhide
Collaborative music agents need internal representations rich enough to support both understanding and generation, yet flexible enough for a workflow where the human retains agency. We present a hierarchical self-supervised ``world model'' for symbolic music: a 2.55M-parameter Swin V2 encoder trained on MIDI piano-roll images with JEPA-style objectives (pitch- and time-shift equivariance, masked embedding prediction, and a distributional regularizer), using no labels and no music-theory vocabulary. Probing the frozen embeddings shows that the level at which a musical property becomes decodable tracks its musical time scale: phrase boundaries are read off the coarsest levels, note density and harmonic detail off the finest. Temporal and phrase structure emerge from the self-supervised objectives alone, while harmonic content must be asked for; a small chord-supervision head raises joint chord recovery from .18 to .54, and key detection, which is never supervised, from .16 to .70. Following the Representation AutoEncoder paradigm, a conditional flow-matching model stands in for a trained decoder, flowing in pixel space from PCA-reduced conditioning: it reproduces a target window at pixel F1 $0.996$, and the same per-level conditioning dropout that controls how far variations stray also enables graphical prompting for masked inpainting with no inpainting-specific sampler. The pipeline runs on CPU producing a suggestion in $2.8$ s, or $0.6$ s on Apple MPS, which we demonstrate in a live interactive demo. In concert with an LLM-based brain, these capabilities supply the core of a collaborative music creation agent in service of, rather than in place of, human agency.
1 Paradigm: “An AI Rick Rubin”
The paper proposes an “AI Rick Rubin” that supports human-led songwriting through careful listening, verbal feedback, and limited musical suggestions rather than taking over creation. Its “ears” and “mouth” combine perceptual representations with controllable generation, under a constraint of timely CPU-only inference.
- Human agency: The “AI Rick Rubin” metaphor denotes a collaborator that uses small models and human expertise to offer decisive feedback while preserving human agency.Rubin’s model emphasizes knowing what he likes and expressing that judgment rather than technical musical ability.
- System scope: The paper develops the agent’s “ears” and “mouth”: perceptual representations drive controllable generation of music suggestions, while an expert-knowledge “Brain” remains out of scope.Figure 1 assigns perceptual representations to Section 3, controllable generation to Section 4, and excludes the reasoning layer.
- Human agency: The system listens to artists’ ideas and returns verbal feedback plus limited musical suggestions, leaving implementation to the human.It is designed as a verbal articulator that “barely” plays instruments.
- Articulation: Description and reproduction require articulable musical structure, unlike mathematical feedback signals such as gradients that human collaborators cannot directly act on.The paper distinguishes communicating what structures are present from steering a system toward an endpoint.
- Representation design: The internal representations use a joint-embedding predictive architecture, or JEPA, following prior applications to musical audio and symbolic music.The cited symbolic-music adaptation found positional information dominated and limited transfer to genre and style content tasks.
- Accessibility: Timely inference on CPU-only runtimes is a fundamental engineering constraint intended to improve accessibility for GPU-poor musicians and avoid sending works-in-progress to cloud services.The constraint was demonstrated through a live demo linked in the supplemental materials.
2 How: A Symbolic Music “World Model”
The model learns hierarchical musical representations without supervised labels or music-theory vocabulary, reflecting musical structure across notes, chords, phrases, and song structure. Its level-wise objective combines shift-aware equivariance, Gaussian regularization, and masked embedding prediction, with selected losses restricted to coarser levels.
- Representation learning: The self-supervised model builds hierarchical representations without relying on labeled data or music-theory vocabulary.It explores the space of musical compositions while learning representations corresponding to notes, chords, phrases, and song structure.
- Training objective: The level-wise training objective combines shift-aware equivariance, sketched isotropic Gaussian regularization, and masked embedding prediction.Equivariance preserves pitch/time translation structure, SIGReg discourages representational collapse, and LMEP predicts EMA-teacher embeddings from student features.
- Training objective: VISReg replaces the Epps–Pulley test with a sliced-Wasserstein distance, but produced no consistent improvement in the reported metrics.The comparison concerns the Gaussian regularization component of the objective.
- Training objective: Equivariance and SIGReg apply only at L0–L3 because SIGReg disrupts fine note-level structure and patch-level equivariance is less meaningful at L4 and L5.The losses are therefore omitted from the two finest levels.
3 Ears: What Does It Actually “Perceive”?
MRJS embeddings reveal musical properties at hierarchy levels aligned with their time scales, while probe results also expose limits of temporal-shift sensitivity and harmony representation. Dataset scale does not consistently improve transfer, and some larger-subset probes degrade.
- Hierarchical perception: Phrase boundaries peak at coarse hierarchy levels L0–L2, whereas note density and harmonic detail are strongest at finer levels.The table reports each probe at the hierarchy level where it scores best, making the winning level informative about musical time scale.
- Probe limitations: .22–.26 temporal-offset scores remain weakest for every model because piano rolls are nearly stationary along time.Pitch is absolute, so vertical shifts change register, whereas sliding a crop sideways often preserves similar material.
- Probe limitations: Harmony scores measure representation capacity rather than downstream performance, since small dedicated chord detectors perform substantially better.The paper frames chord and key probes as capacity tests, not substitutes for specialized downstream detectors.
- Dataset variation: 4× Lakh training data produced no consistent gain over 1× at matched budget and degraded chroma (.76 to .61) and cross-song separation (.62 to .77).Both subsets transferred well on POP909, matching or exceeding MRJS on note density and cross-song separation.
4 Mouth: Making Suggestions
The section realizes musical suggestion and inpainting with a pixel-space flow-matching model conditioned on PCA-reduced embeddings of the user’s existing idea. Per-level conditioning dropout supports user-controlled variation and masked inpainting, while imperfect fits frame outputs as suggestions for human implementation.
- Conditioned generation: Pixel-space flow matching generates musical windows from noise, conditioned on embeddings of the user’s existing idea.The model supports replacing melodies or accompaniments and suggesting different continuations.
- Conditioned generation: F1=0.996: PCA-reduced conditioning reconstructs the input window closely despite reducing the signal by over 3×.PCA keeps at least 90% of variance per level, reducing 16,128 to 4,835 floats.
- Inpainting: Per-level conditioning dropout enables masked inpainting and ordinary generation, with masked patches removed under user-chosen strengths and untouched input preserved exactly.The pipeline uses 10 Euler steps with guidance 1; without a mask, it performs ordinary conditioned generation.
- Human agency: Infilled notes generally fit surrounding music but vary in commitment and sometimes land off, making suggestions material for the user to implement rather than finished replacements.The approach is framed as an AI Rick Rubin that “barely” plays the instrument.
5 Outro: What To Do With This?
A live demo makes the model’s representational control tangible while meeting the stated timely-interaction target on CPU. Encoding is fast, and suggestions complete in seconds, with Apple’s Metal backend slower for encoding.
- The live demo serves as both an interactive probe of representational capacity and an application showing how the model’s control feels in practice.The paper describes the app as a firsthand demonstration of interactive control.
- 8.6 ms encodes a 128 × 128 window on two CPU threads, faster than Apple’s Metal backend at 19.9 ms.These timings are reported for the M1 Max setup.
- 2.8 s produces a suggestion on a full CPU, compared with 3.8 s on two CPU threads, meeting the paper’s “timely” design target.The suggestion uses 10 Euler steps at guidance 1.0 and one function evaluation per step.
A Encoder training objective: equations and details
The encoder objective combines four level-specific self-supervised terms, using shifted 128 × 128 piano-roll views, smooth distance-based equivariance, masked embedding prediction, distribution regularization, and soft pitch/time factorization. Loss placement reflects hierarchy: equivariance and SIGReg operate at L0–L3, while factorization operates only at L0–L2.
- Objective structure: Four terms comprise the total encoder training objective, with losses assigned selectively across hierarchy levels.Lequiv and LSIGReg apply at L0–L3, while Lfact applies at L0–L2.
- View generation: Two full 128 × 128 views are generated from each piano roll by sampling a second crop after random time- and pitch-axis shifts.The original crop is x1, while x2 is offset by (∆x, ∆y) pixels and retains the encoder input size.
- Equivariance loss: The equivariance loss targets embedding distances proportional to normalized shift magnitude, avoiding collapse of shifted pairs to identical representations.The smooth quadratic both attracts pairs that are too distant and repels pairs that are too close; shifts use Beta(2, 2), with ∆pmax = 12 and ∆tmax = 48.
- Distributional regularization: SIGReg enforces an isotropic Gaussian embedding prior, and chunking 256 slices into groups of 32 reduces peak VRAM by ∼5 GB.SIGReg is applied only to student embeddings z2 and accumulated in float32.
- Masked embedding prediction: An EMA teacher with η = 0.96 supplies stable targets for equivariance and masked embedding prediction, while a lightweight predictor reconstructs embeddings at masked patch positions.The teacher processes x2 without masking, and prediction uses the student’s full-context representation.
- Soft factorization: Soft factorization constrains augmentation-difference geometry so same-sign shifts are parallel, opposite-sign shifts anti-parallel, and cross-type shifts orthogonal.The cosine-similarity loss is applied only at L0–L2, without prescribing the directions themselves.
B Full STORMBIRD probe results
Appendix B reports the complete per-level STORMBIRD3 probe suite across all trained encoders, complementing the main text’s best-level summary. The results cover musical, structural, temporal, and affective representation metrics with stated chance levels and optimization directions.
- Probe-suite scope: STORMBIRD3 reports per-level results for the complete probe suite across all trained encoders, while the main text reports each model’s best level.Values are reported to 3 decimals.
- Level-wise factorization: Soft factorization results are shown for levels L1–L5, although factorization loss is applied only at L0–L2.Geometry at L3–L5 is inherited through the shared backbone rather than directly enforced.
- Musical content probes: The appendix evaluates chord quality accuracy, root-note accuracy, chroma regression R2, and 24-class key-detection accuracy.The root-note task is 12-class with chance .083; chord quality has chance ≈.50.
- Structural and temporal probes: The suite measures note density R2, cross-song distance ratio, and time-translation R2 to assess density, song-level separation, and temporal representation.Cross-song distance ratio is lower-is-better; the other listed metrics are higher-is-better.
C Equivariance
This section evaluates equivariance through per-level embedding distances under pitch transposition and time translation, using R2 to measure how predictably distances track shift magnitude.
- Pitch transposition: Pitch-transposition equivariance is assessed by measuring per-level embedding distance across transposition shifts.Figure 9 reports results for chall-scr1.5, with shaded bands showing ±1 standard deviation over probe samples.
- Pitch transposition: Pitch-transposition equivariance is summarized with an R2 metric, where higher values indicate more predictable ordering of distances by shift size.The R2 values come from linear fits to the probes’ logged distance curves.
- Time translation: Time-translation equivariance is summarized using a corresponding R2 metric.The section reports this evaluation in Table A13.
D Phrase-boundary analyses
Phrase-boundary understanding is evaluated against human annotations using unsupervised novelty segmentation and linear boundary-detection probes across encoder levels. The analyses report explicit chance, fixed-grid, raw-pixel, and human-ceiling reference points for interpreting performance.
- Evaluation setup: Two measurements use Dai et al.’s human phrase annotations: unsupervised Footenovelty segmentation F1 and linear boundary-detection probes by encoder level.The probe reports average precision and AUC, while segmentation uses tolerance in 32nd-note-resolution columns.
- Unsupervised segmentation: .275 is the chance F1@32 for count-matched random phrase-boundary placement, compared with .405 for a fixed 8-bar grid.Both references apply to the unsupervised Footenovelty segmentation measurement.
- Unsupervised segmentation: .895 is the human inter-annotator ceiling for the phrase-boundary segmentation evaluation.This ceiling contextualizes the unsupervised segmentation score.
- Linear probing: .17 is the AP chance level and .50 is the AUC chance level for the linear boundary-detection probe, while raw 32×32 pixels score .18/.51.The raw-pixel baseline is reported as AP/AUC, respectively.
E Ablation studies
The ablation studies screen depth allocation, per-level SIGReg strength, and chord-supervision weight across controlled training variants. They track content probes, cross-song performance, and catastrophic L5 chroma collapse, with validation losses compared only within each table.
- Screening protocol: Screening results mark catastrophic L5 chroma collapse with ×, including R2 values from −77 to −713 despite normal training curves.Missing data are marked —, and validation losses are comparable only within a table.
- Depth allocation: Depth-allocation screens evaluate content probes at L5 and cross-song performance as the minimum across levels.The campaign uses 100-epoch screens; variants with catastrophic L5 chroma collapse are marked separately.
- SIGReg strength: Per-level SIGReg ablations vary the λ schedule on the uniform3 backbone.The reference uses scalar λ=0.15 for L0–L3 only.
- Chord supervision: Chord-supervision ablations vary λchord for heads on L0–L3 or all levels.“/all” variants are 100-epoch screens, while the remaining variants run for 250 epochs; values report the best level.
F Runtime
The runtime evaluation measures the complete encode–sample–render demo pipeline across backends using 10 Euler steps and classifier-free guidance. Sampling cost depends on guidance strength: at exactly 1.0, the unconditional branch is skipped and sampling times halve.
- F Runtime: The benchmark measures wall-clock time for the full encode–sample–render demo pipeline across backends.Parenthesized values represent speedups over the 2-thread CPU baseline.
- F Runtime: 10 Euler steps require 20 function evaluations under classifier-free guidance.This is the general-case configuration used for the reported timings.
- F Runtime: At guidance strength exactly 1.0, skipping the unconditional branch halves sampling time.The general case evaluates both guidance branches at each Euler step.