AI Choreographer: Music Conditioned 3D Dance Generation with AIST++

In plain words

Picture a dancer who has just moved for two seconds and then hears a song. The model's job is to keep dancing: it looks at the recent moves and the music and writes out what comes next, one frame at a time, so the motion stays in time with the beat. Earlier systems trained on this kind of data often froze after a few seconds, drifted away, or hardly used the music. The authors' model, FACT, addresses this with three choices. Every frame may look at every other frame. It is trained to predict a block of upcoming frames, not just the next one. The music and the motion are mixed early, inside one deep shared network.

Training needed data that did not exist yet. The authors rebuilt 3D dance motion from existing multi-view dance videos, paired it with the music, and called the result AIST++. The reported results are the authors' own, compared against a few earlier methods.

TL;DR

• Generates 3D dance motion (joint rotations plus global translation) from music and a 2-second seed motion. The model is FACT §4.1.
• Three design choices carry the result: full attention, predicting N = 20 future frames, and early fusion in a 12-layer cross-modal transformer §1 tab. 3 tab. 4.
• Ships with the AIST++ dataset: 1408 sequences, 5.2 hours, 10 genres, nine camera views per frame §3 tab. 1.
• Author-reported on the AIST++ test split: FID_k 35.35 and FID_m 12.40, both better than the three baselines; BeatAlign 0.241 tab. 2.
• The released code differs from the paper in batch size, learning-rate steps and training length. See Paper vs code before reproducing numbers.

Problem

• Dance must be continuous, kinematically complex and synchronized with music, and the mapping from music to motion is non-linear §1.
• Motion capture of dancers needs heavy instrumentation, so earlier 3D dance datasets were small §1. The AIST source is multi-view video without camera calibration, so 3D motion has to be reconstructed first §3.
• A causal transformer trained to predict only the next frame, as used in earlier dance and motion work, freezes after several seconds. A full-attention model trained on one future frame drifts §4.1 §5.2.3 tab. 3.
• One choreography appears with several music tracks at different tempos. A model can therefore learn to ignore the music unless the audio branch is fused deeply §3 §4.1.
• The model is kinematic. It does not reason about floor contact, so global translation can produce foot sliding or floating §6.

Method

1. Dataset AIST++ §3 app. A. Start from the AIST Dance Database. Cameras are calibrated by 2-D pose detection and bundle adjustment. Per frame, 17 COCO joints are triangulated with temporal-smoothness and bone-length terms. SMPL is then fitted with the mean shape fixed app. A. Each frame carries nine views of camera parameters, 17 COCO joints in 2-D and 3-D, and 24 SMPL poses with global scale and translation §3. Split §5.2.1: the test set holds one music piece per genre, two dancers, and two choreographies per piece, giving 40 choreographies. The training set holds 329 choreographies. No music or choreography is shared between the two.

2. Features §5.2.1.

   • Motion: 219 values per frame. Each of 24 joints gives a 3x3 rotation matrix (9 values), plus a 3-value global translation.
   • Audio: 35 values per frame at 60 FPS, computed with librosa. The parts are the onset envelope (1), 20 MFCCs, 12 chroma bins, a one-hot onset-peak flag (1) and a one-hot beat flag (1) preprocessing.py.

3. FACT architecture §4.1 §5.2.1.

   1. Project motion (219 to 800) and audio (35 to 800) linearly. Add a learned positional embedding to each stream.
   2. Run a motion transformer (2 layers) and an audio transformer (2 layers). Both use 10 heads and width 800.
   3. Concatenate the motion tokens and then the audio tokens into one sequence.
   4. Run a cross-modal transformer (12 layers, 10 heads, width 800). It has no positional encoding. A linear layer maps its output to 225 values.

   No attention mask is used anywhere, so every token attends to every input token §4.1. The paper writes the attention block as:

   \mathbf{C} = \text{FF}(\text{Attn}(\mathbf{Q,K,V,M})) = \text{FF}\Big(\text{softmax}\Big(\frac{\mathbf{QK}^T+\mathbf{M}}{\sqrt{D}}\Big)\mathbf{V}\Big), \quad \mathbf{Q}=\mathbf{XW}^Q,\ \mathbf{K}=\mathbf{XW}^K,\ \mathbf{V}=\mathbf{XW}^V \qquad [\text{eq. 1}]

   Here M is the mask, which FACT leaves empty, and D is the number of channels in the attention layer §4.1. The paper compares the early and late fusion layouts, and a single MLP fusion used by earlier work §4.1 §5.2.3.

4. Future-N supervision §4.1 §5.2.1. The model reads a 120-frame seed and is trained to output the next N = 20 frames. The loss is the mean squared error over those frames. Only the first N output positions are compared with the targets fact_model.py.

5. Training §5.2.1. Adam, initial learning rate 1e-4, batch size 16 (see Hyperparameters). Each training clip gives one random window, sampled as in the pseudocode below.

# Training window sampling  [code: mint/utils/inputs_util.py#L74-L105]
# Input:  motion (T, 225) with 6 leading zeros; audio (T', 35) at 60 FPS
# Output: motion_input (120, 225), target (20, 225), audio_input (240, 35)

def sample_training_window(motion, audio):
    # Window length: max(120, 120 + 20, 240) frames  [§5.2.1]
    window = 240
    # Random start frame inside the clip  [code: inputs_util.py#L79-L84]
    start = random_integer(0, length(motion) - window + 1)
    # Seed: first 120 frames of the window  [§5.2.1]
    motion_input = motion[start : start + 120]
    # Target: the 20 frames right after the seed  [§4.1]
    target = motion[start + 120 : start + 140]
    # Music: 240 frames from the same start, aligned to the seed  [§5.2.1]
    audio_input = audio[start : start + 240]
    return motion_input, target, audio_input

# Future-N loss  [§4.1] [code: mint/core/fact_model.py#L143-L148]
# Input:  predicted (360, 225) from the cross-modal output; target (20, 225)
# Output: scalar

def future_loss(predicted, target):
    # Keep only the first N output positions  [code: fact_model.py#L146]
    difference = target - predicted[0 : length(target)]
    # Mean over all values, incl. 6 zero dims  [code: fact_model.py#L147]
    return mean(square(difference))

6. Generation §5.2.1 §5.2.2 fact_model.py. Generation is auto-regressive. Only the first predicted frame is kept each step, and both the motion window and the music window slide by one frame. The evaluator runs 1200 steps, which is 20 seconds at 60 FPS [code: mint/ctl/single_task_evaluator.py#L69].

# Auto-regressive generation  [§5.2.1] [code: fact_model.py#L120-L132]
# Input:  seed (120, 225) motion; audio (T', 35) at 60 FPS
# Output: frames (steps, 225) of generated motion

def generate(model, seed, audio, steps=1200):
    # The motion window starts as the seed  [§5.2.1]
    motion_window = seed
    generated = []
    for step in range(steps):
        # Music window of 240 frames, moved forward one frame per step  [§5.2.1]
        audio_window = audio[step : step + 240]
        # One full forward pass over 120 + 240 tokens  [§4.1]
        predicted = model(motion_window, audio_window)  # (360, 225)
        # Keep only the first predicted frame  [code: fact_model.py#L128]
        next_frame = predicted[0 : 1]
        generated.append(next_frame)
        # Slide the motion window: drop the oldest frame, append the new one
        motion_window = concat(motion_window[1 : ], next_frame)
    return concat(generated)

7. Evaluation §5.2.2 §5.2.4.

   • FID: Frechet distance between Gaussian fits of real and generated motion features. The paper uses two feature sets, a geometric one (33 boolean features) and a kinetic one (72 values). Each model contributes 40 generated sequences of 1200 frames §5.2.2.
   • Diversity: average Euclidean distance in feature space across the 40 generated motions, in each feature space §5.2.2.
   • BeatAlign: for each kinematic beat, the closest music beat is found. The score averages exp(-d^2 / (2 sigma^2)) over kinematic beats eq. 2. The paper sets sigma = 3 for 60 FPS. Kinematic beats are local minima of the kinetic velocity §5.2.2.
   • User study: 30 participants per baseline. Each watches 10 pairs of FACT and a random counterpart and picks the one that dances more to the music §5.2.4.

   \text{BeatAlign} = \frac{1}{m}\sum_{i=1}^{m} \exp\Big(-\frac{\min_{t_j^y \in B^y}\lVert t_i^x - t_j^y\rVert^2}{2\sigma^2}\Big) \qquad [\text{eq. 2}]

Hyperparameters

Results

Author-reported. Setup §5.2.1 §5.2.2: AIST++ test split (40 choreographies). Baselines retrained on the same data with the authors' official code. Each model generates 40 sequences of 1200 frames. FID is computed against the AIST++ test motions with the paper's feature extractors. Diversity and BeatAlign are the paper's own metrics.

(*) The paper notes Li et al.'s output is discontinuous, which inflates its kinetic FID tab. 2.

User study, FACT win rate in pairwise comparisons tab. 2 §5.2.4: vs. Li et al. 80.6%; vs. DanceNet 71.1%; vs. DanceRevolution 77.0%; vs. random AIST++ motion 25.4% (so the random clip was preferred about 75% of the time).

Ablation on attention and supervision tab. 3:

Ablation on fusion tab. 4:

(**) Computed with the music paired with the input motion tab. 4 caption.

Paper vs code

Code: google-research/mint at b8f8bdfbbe3fbfa67831a2ef9bcf71a6a9e74552, Apache-2.0. The paper's project page links this repository as the model code.

• Batch size. Paper: 16 §5.2.1. Config: 32 fact_v5_deeper_t10_cm12.config.
• Learning-rate schedule. Paper: drops after 60k and 100k steps §5.2.1. Config: drops after 100k and 150k steps fact_v5_deeper_t10_cm12.config.
• Training length. Paper: finishes after 300k steps §5.2.1. The trainer default is 2,400,000 steps, and the config does not set a value trainer.py trainer.py.
• Config name. The file is named t10, but the target length is 20 frames fact_v5_deeper_t10_cm12.config. The paper uses N = 20 §4.1. The config name is not explained in the repository.
• Motion translation scale. The paper does not mention it. The code divides translation by a per-clip SMPL scale factor before use preprocessing.py preprocessing.py.
• Zero padding. The paper gives 219 motion values §5.2.1. The model reads 225 values, with six zeros prepended to the front inputs_util.py. The evaluator reads translation from indices 6 to 8 and rotations from index 9 on calculate_scores.py.
• Rotation projection at evaluation. The paper does not mention it. Predicted 3x3 rotation matrices are projected to the nearest proper rotation by SVD before conversion to axis-angle calculate_scores.py calculate_scores.py.
• Context length. The paper calls the context window 240 §4.1. The cross-modal sequence is 120 motion tokens plus 240 audio tokens, so 360 tokens base_models.py fact_v5_deeper_t10_cm12.config fact_v5_deeper_t10_cm12.config.
• Feature extractors for FID. The paper uses fairmotion extractors §5.2.2. The code's comment says the paper's manual (geometric) extractor came from a buggy fairmotion version. It also reports FID_k about 38 and FID_g about 27 calculate_scores.py, which do not match the paper's 35.35 and 12.40.
• Reference set and normalisation. The paper says the real motions are the AIST++ test set §5.2.2. The code loads every cached file in data/aist_features calculate_scores.py and normalises both sets with the real set's mean and standard deviation calculate_scores.py. The paper does not state either choice.
• Test set construction. The code's test split adds unpaired music, reusing sequences ten times with randomly chosen music preprocessing.py. The paper reports 40 choreographies §5.2.1. The split files are not in the repository.
• Metrics not in the code. The repository has no BeatAlign, diversity or user-study code. A search for beat, align and diversity in the Python files found nothing relevant.
• Released configs. Only the main configuration is included. The future-1, future-10, causal and fusion ablations have no config here .
• Dropout. The paper does not mention dropout. The model code path applies none. The proto default of 0.1 is never used by the transformer blocks model.proto.
• Warmup. The paper does not mention warmup. The trainer's warmup flag applies only to exponential and cosine schedules, not to the manual schedule used here trainer.py train.proto.

No discrepancies were found in the pieces that match: the 120-frame seed, the 240-frame music window, the 20-frame target, the 800 width, 10 heads, the 2 + 2 + 12 layer layout, the 35-value audio feature, and 1200-step generation §5.2.1 §5.2.2.

Pitfalls

• Attention scale. The code scales dot products by 1/sqrt(800), using the full hidden width rather than the per-head width of 80 base_models.py. The paper's D is only described as the channel count of the attention layer §4.1. If you reproduce the numbers, follow the code.
• Only the first output frame is used. Generation keeps one predicted frame per step fact_model.py. The other 19 supervised frames act only as training signal fact_model.py. Because the supervised output positions are indexed from the start of the output block, not aligned to input positions, the first output position must learn to mean the next frame.
• Full recomputation each step. The rollout calls the whole model on 360 tokens at every step, with no caching across steps fact_model.py. Generating 1200 frames therefore means 1200 full forward passes.
• Loss scale. The loss is averaged over all elements, including the six always-zero leading dims fact_model.py. Changing the feature layout changes the effective scale of the loss, so keep the 225-value layout when matching the paper's optimisation.
• Only N = 20 is released. The ablation table shows future-10 with a lower FID_k than future-20, 35.10 versus 35.35 tab. 3. The paper's choice rests on FID_m and BeatAlign. Treat the N = 20 advantage as a multi-metric trade-off.
• FID numbers are feature-dependent. The paper says earlier motion encoders were not public, so it uses fairmotion features §5.2.2. The repository comment says the geometric features used in the paper had a bug, and its printed FID_g does not match calculate_scores.py. Do not compare FID across papers that use other extractors.
• Normalisation changes FID. Both feature sets are normalised with the real set's statistics calculate_scores.py. Using a different normalisation gives different values.
• BeatAlign has an undefined input. The paper defines kinematic beats as local minima of kinetic velocity §5.2.2, but no velocity definition or beat-detection code is in the repository. The paper says sigma = 3 at 60 FPS. Implementers must choose the velocity definition themselves.
• Internal inconsistencies. PCKh@0.5 is 98.7% in §5.1 but 98.4% in the appendix §5.1 app. A. FACT's FID_m is 12.40 in the main table and 12.39 in the ablation tables tab. 2 tab. 3 tab. 4. Quote the source table explicitly.
• Win-rate column. The column is headed "FACT WinRate" with a down arrow, but the values are FACT's win rates. Higher means FACT was preferred tab. 2.
• User study ceiling. Real, highly expressive AIST++ motion was preferred about 75% of the time against FACT. The authors attribute this partly to participants preferring that expressiveness §5.2.4.
• Determinism. The model is deterministic, so one seed produces one output. Diversity comes from different music and seeds, not from sampling §6.
• Kinematic output. Global translation can cause foot sliding and floating, because no contact or physics is modelled §6.
• Compute. Training took about 3 days on 4 TPUs §5.2.1. The authors describe full attention as feasible because the context is short §4.1.
• Dataset licence. The paper does not state a dataset licence. The project page says the assets are CC BY 4.0. Check the dataset's own terms before reuse.

Lineage

• Builds on: the Transformer encoder, used with full attention, layer-norm and GELU feed-forward blocks §4.1 base_models.py. Causal GPT-style decoding is the contrast the paper draws §4.1.
• Compared against: Li et al. (2020), DanceNet (Zhuang et al., 2020) and DanceRevolution (Huang et al., 2021), all retrained with official code tab. 2 §5.2.1.
• Prior causal, shift-by-1 training [Aksan et al. 2020, cited in §4.1] is the approach the future-N supervision replaces.
• Superseded by: not known from the content of this wiki.