MoMask: Generative Masked Modeling of 3D Human Motions

In plain words

Computers can now write text and draw pictures, but making a realistic 3D character move from a written sentence is harder. One common approach turns motion into a short list of discrete "motion words" and then predicts those words one at a time, the way a language model predicts text.

MoMask improves this in two ways. First, each motion is stored as a stack of word layers. The first layer captures the main shape of the movement, and each layer above it stores what the layers below missed, so the stack gets closer and closer to the original motion. Second, the model fills in the first layer by hiding some of its words and guessing them together from the text and the words it can already see. It does this in parallel and in a few rounds rather than word by word. Once that layer is complete, a second model fills in the finer layers.

This matters because the method is both more accurate and faster than word-by-word generation, and the same model can also fill in a gap inside an existing motion, for example to change one part of a clip and keep the rest.

TL;DR

MoMask is a text-to-motion generator for 3D human motion. It uses residual vector quantization to represent each motion as six token layers, a masked transformer for the base layer and a residual transformer for the others. The authors report state-of-the-art FID, R-Precision and multimodal distance on HumanML3D and most of these on KIT-ML tab. 1. Relevant if you need text-driven motion with few sampling steps, or motion inpainting.

Problem

Earlier discrete motion generators quantize each motion once, so the quantization error caps the reconstruction quality §1. Autoregressive decoding predicts motion tokens in one direction, so each token sees only its past and errors accumulate §1. Discrete diffusion alternatives decode in both directions but need hundreds of iterations §1. The paper assumes text conditioning and a motion-language dataset with text for each motion §4.

Method

Overview. A residual VQ-VAE turns a motion into V+1 token layers. A masked transformer models the base layer, and a residual transformer models the other V layers §3 fig. 2.

Residual VQ-VAE §3.1. A 1D convolutional encoder maps a motion of N frames with D features to n latent vectors of width d, with downsampling ratio n/N §3.1. Each quantization layer replaces its input residual with the nearest codebook entry, and the next residual is the difference §3.1 eq. 1. The final approximation is the sum of all layer codes, which the decoder maps back to a motion §3.1. Training minimizes an L1 reconstruction loss plus an embedding term summed over the residual layers v = 1..V §3.1 eq. 2. Codebooks are updated by exponential moving average with codebook reset, following T2M-GPT, and gradients pass through the quantizer with a straight-through estimator §3.1.

Quantization dropout §3.1. During training, with probability q the last k layers are disabled, where k is between 0 and V, so the first layers must recover as much of the motion as possible §3.1.

Masked transformer §3.2. A bidirectional transformer models the base-layer tokens. Text features come from a CLIP text encoder §3.2. For each sequence a fraction of positions is replaced with a mask token. The fraction follows a cosine schedule, gamma(tau) = cos(pi * tau / 2) with tau drawn uniformly from [0, 1] §3.2. The loss is the negative log-likelihood of the original tokens at masked positions §3.2 eq. 3. BERT-style replacement applies to selected masked positions: 80% mask token, 10% random token, 10% unchanged §3.2.

Residual transformer §3.3. A second bidirectional transformer predicts the residual layers. It has one embedding table per residual layer. Embeddings of all layers before the target layer j are summed and fed in with the text and the index j §3.3. The layer j tokens are predicted in parallel §3.3 eq. 4. The projection of layer j and the embedding of layer j+1 share weights §3.3.

Inference §3.4 fig. 3. Starting with all base tokens masked, the masked transformer runs for L iterations. At each iteration it samples the masked positions, and the sampled tokens with the lowest confidence are masked again, following the cosine schedule §3.4. After the base layer is complete, the residual transformer predicts layers 2..V one after another §3.4. The RVQ decoder then maps the tokens to a motion §3.4.

Classifier-free guidance §3.4. Both transformers are trained with the condition dropped with probability 10% §3.4. At inference the guided logits are formed at the final linear layer, before softmax, from the conditional logits omega_c and the unconditional logits omega_u with scale s §3.4 eq. 5.

Temporal inpainting §4.3 fig. 6. Tokens in a chosen time range are masked and the same inference procedure runs, which regenerates the masked range conditioned on text and keeps the rest of the sequence §4.3.

Key equations.

\mathbf{b}^v = \mathrm{Q}(\mathbf{r}^{v}), \quad \mathbf{r}^{v+1} = \mathbf{r}^{v} - \mathbf{b}^{v}, \quad v = 0,\dots,V eq. 1

\mathcal{L}_{rvq} = \|\mathbf{m}-\mathbf{\hat{m}}\|_1 + \beta\sum_{v=1}^V\|\mathbf{r}^{v}-\mathrm{sg}[\mathbf{b}^v]\|_2^2 eq. 2

\mathcal{L}_{mask} = \sum_{\tilde{t}_k^0=[\mathrm{MASK}]}-\log p_\theta(t_k^0| \tilde{t}^0, c) eq. 3

\mathcal{L}_{res} = \sum_{j=1}^V\sum_{i=1}^n-\log p_\phi(t_i^j| t_i^{1:j-1}, c, j) eq. 4

\omega_g = (1+s)\cdot\omega_c - s\cdot\omega_u eq. 5

Training and inference, pseudocode. Layer 0 is the base layer and layers 1 to V are residual. These sketches follow the paper's text, not a line-by-line port of the code. Where the code differs, see Paper vs code.

# Residual quantization of one latent sequence  [§3.1] [eq. 1]
# Input:  latents (B, n, d) from the 1D conv encoder  [§3.1]
# Output: summed quantized latents, one token sequence per layer

def residual_quantize(latents, codebooks, num_layers):
    # Start from the encoder output as the residual r^0  [§3.1]
    residual = latents                             # (B, n, d)
    quantized_sum = zeros_like(latents)            # (B, n, d)
    token_layers = []
    for layer in range(num_layers):
        # Nearest code for the current residual  [§3.1]
        indices = nearest_code(residual, codebooks[layer])
        code = codebooks[layer][indices]           # (B, n, d)
        # The next residual is what this layer missed  [§3.1] [eq. 1]
        residual = residual - code
        quantized_sum = quantized_sum + code
        token_layers.append(indices)
    return quantized_sum, token_layers

# RVQ-VAE training loss  [§3.1] [eq. 2]
# The paper leaves the commitment weight beta unstated.

def rvq_loss(motion, recon, residuals, codes, beta):
    # Reconstruction term: L1 between motion and decoder output  [§3.1]
    rec = mean(abs(motion - recon))
    # Embedding term: pull each residual toward its code, codes fixed  [§3.1]
    commit = 0.0
    for v in range(1, num_layers):
        commit = commit + sum_squares(residuals[v] - stop_grad(codes[v]))
    return rec + beta * commit

# Quantization dropout  [§3.1]
# Returns how many layers to keep for this training step.

def layers_to_keep(num_layers, q):
    # With probability 1 - q keep every layer  [§3.1]
    if random() >= q:
        return num_layers
    # Otherwise disable the last k layers, k between 0 and V  [§3.1]
    return randint(1, num_layers)

# One masked-transformer training step  [§3.2] [eq. 3]
# tokens: (B, n) base-layer indices; text: (B,) captions

def masked_training_step(tokens, text):
    # Text condition from a frozen CLIP encoder  [§3.2]
    cond = clip_text_encoder(text)                 # (B, d)
    # Draw a mask ratio from the cosine schedule  [§3.2]
    tau = uniform(0, 1)                            # (B,)
    ratio = cos(pi * tau / 2)                      # (B,)
    count = ceil(ratio * n)                        # [§3.2] (code rounds)
    # Choose that many random positions to mask  [§3.2]
    masked = random_positions(n, count)            # (B, n) bool
    # Of the masked positions: 80% mask token, 10% random, 10% kept  [§3.2]
    inputs = bert_replace(tokens, masked)
    # Negative log-likelihood of the original tokens at masked positions  [§3.2]
    logits = masked_transformer(inputs, cond)      # (B, n, K)
    return cross_entropy(logits[masked], tokens[masked])

# One residual-transformer training step  [§3.3] [eq. 4]
# indices: (B, n, V+1) layer tokens; text: (B,) captions

def residual_training_step(indices, text):
    # Pick one residual layer j to learn  [§3.3]
    layer = choose_residual_layer(num_layers)      # j in [1, V]
    # Sum embeddings of all layers before j  [§3.3]
    history = 0.0
    for lower in range(layer):
        history = history + token_embed[lower](indices[..., lower])
    # Predict layer j tokens in parallel  [§3.3] [eq. 4]
    logits = residual_transformer(history, text, layer)
    return cross_entropy(logits, indices[..., layer])

# Iterative base-layer decoding  [§3.4]
# The paper runs L = 10 iterations for MoMask  [§4]

def generate_base_tokens(text, length, steps):
    # Start with every base-layer position masked  [§3.4]
    tokens = full(length, MASK_ID)                 # (n,)
    for step in range(steps):
        # Predict all masked positions at once  [§3.4]
        logits = masked_transformer_cfg(tokens, text)     # (n, K)
        sampled = sample_categorical(logits)               # (n,)
        # Keep the new samples only at masked positions  [§3.4]
        is_masked = tokens == MASK_ID
        tokens = where(is_masked, sampled, tokens)
        # Confidence of each new sample; fixed tokens are never re-masked  [§3.4]
        confidence = where(is_masked, prob_of(logits, sampled), INFINITY)
        # Re-mask the least confident tokens for the next round  [§3.4]
        count = ceil(cos(pi * (step + 1) / (2 * steps)) * length)
        lowest = argsort(confidence)[:count]
        tokens[lowest] = MASK_ID
    return tokens

# Residual layers after the base layer  [§3.4] [§3.3]
# base_tokens: (n,) from generate_base_tokens

def generate_residual_layers(base_tokens, text, num_layers):
    # Fill layers 1..V in order, each conditioned on the layers before it  [§3.4]
    layers = [base_tokens]
    for layer in range(1, num_layers):
        logits = residual_transformer_cfg(layers, text, layer)
        layers.append(sample_categorical(logits))
    # The RVQ decoder turns the stacked tokens into a motion  [§3.4]
    return layers

# Classifier-free guidance on logits  [§3.4] [eq. 5]
# Paper form, with s as written in eq. 5

def guided_logits(cond_logits, uncond_logits, scale):
    # (1 + s) * conditional - s * unconditional  [§3.4] [eq. 5]
    return (1 + scale) * cond_logits - scale * uncond_logits

Hyperparameters

Paper values are the ones written in the text §4. Code values come from the official repository at code_commit. Code paths are relative to the repository root, and the cited lines are where each default is set.

Results

Author-reported. Setup: HumanML3D (14,616 motions, 44,970 descriptions) and KIT-ML (3,911 motions, 6,278 descriptions), with mirrored augmentation, and evaluation with the T2M metrics over 20 repeats §4. The paper compares against TM2T, T2M, MDM, MLD, MotionDiffuse, T2M-GPT and ReMoDiffuse §4. "MoMask (base)" uses only base-layer tokens §4. Values are mean ± 95% confidence interval. The source marks best and second-best values in bold and underline, which are not reproduced here.

Table 1: text-to-motion generation tab. 1

The paper says ReMoDiffuse leads on KIT-ML FID because it retrieves from a large database §4. Its user study on HumanML3D found MoMask preferred in most comparisons, and 42% preference against ground truth §4 fig. 5. Inpainting was preferred over MDM in 68% of 40 samples judged by 6 users §4.3.

Table 2: reconstruction and ablations tab. 2. MPJPE is in millimetres.

The paper reports that residual layers help reconstruction throughout, and that generation quality starts to degrade above five residual layers §4.2. Guidance scale s near 4 is the sweet spot for the masked transformer, and about 10 iterations already reach the best FID and multimodality distance §4 fig. 7.

Paper vs code

Repository: https://github.com/EricGuo5513/momask-codes at commit 94a6636c9c463b7a9414c3401a6f1b67e6c51824, MIT licensed. It is linked from the authors' project page §1 author block. Reported numbers were not re-run. Paths below are relative to the repository root.

• Transformer depth. Paper: 6 layers §4. Code: default 8 base_option.py, and the README training commands do not override it, while the model names say nlayer8 README.md. The published checkpoints presumably used 8, but this is inferred from the defaults and names, not confirmed.
• Batch size for KIT-ML transformers. Paper: 32 §4. README: 16 README.md.
• Residual-transformer condition drop. Paper: 10% §3.4. README residual command: 0.2 README.md README.md. Masked transformer default 0.1 train_option.py.
• Residual layer selection. Paper: layer j is picked at random §3.3. Code: draws the layer through a cosine-shaped schedule, not uniformly transformer.py transformer.py.
• Guidance formula. Paper: (1+s)*cond - s*uncond eq. 5. Code: uncond + s*(cond - uncond), applied to logits transformer.py. The two agree only if the paper's s is the code's s minus 1. The published values (4, 5, 2) match the code's argument, so the code convention is the one that reproduces the reported numbers. Scale 1 in the code returns the conditional logits only transformer.py.
• Commitment loss. Paper: beta times a sum over layers v = 1..V eq. 2. Code: MSE per layer, averaged over the kept layers, including layer 0 residual_vq.py quantizer.py, weighted by 0.02 vq_option.py vq_trainer.py.
• Reconstruction loss. Paper: L1 eq. 2 text, §3.1. Code: Smooth-L1 over all features, plus 0.5 times Smooth-L1 over the local joint positions vq_trainer.py vq_trainer.py vq_option.py.
• Gumbel noise in RVQ training. Paper: not mentioned. Code: the RVQ forward runs with temperature 0.5 and Gumbel noise during training, so code assignments are stochastic model.py quantizer.py quantizer.py.
• Quantization dropout granularity. Paper: "randomly disables" layers with probability q §3.1. Code: one decision per batch, then a uniform choice of how many layers to keep residual_vq.py. Dropped layers contribute no commitment loss residual_vq.py.
• Mask-ratio rounding. Paper: ceil(gamma * n) §3.2. Code: round with a minimum of 1 transformer.py transformer.py.
• BERT-style replacement. Paper: 80/10/10 §3.2. Code: 10% random tokens, then 88% of the remaining masked positions become mask tokens, so about 79% mask, 10% random and 11% unchanged transformer.py.
• Inference schedule. Paper: gamma(l/L) with l from iteration counts §3.4. Code: linspace(0, 1, steps) as the schedule input, so the step index is l/(L-1), and the last step still re-masks at least one token transformer.py transformer.py.
• Top-k filtering and temperature. Not in the paper §3.4. Code: keeps the top 10% of the 512 logits (ceil(0.1 * vocab)) with temperature 1 transformer.py tools.py eval_option.py transformer.py.
• Confidence for re-masking. Paper: "confidence" of sampled tokens, unspecified §3.4. Code: probability of the sampled token from the untempered, unfiltered guided softmax transformer.py.
• Iteration count. Paper: L = 10 §4. README eval: 10 README.md. Eval default: 18 eval_option.py.
• Residual guidance in evaluation. Paper: 5 for residual, both datasets §4. Code: hard-coded to 5 in the eval helper eval_t2m.py and in the demo gen_t2m.py. Not a CLI option. Editing uses 2 edit_t2m.py; the paper gives no editing value §4.3.
• Optimizer. Paper: not stated. Code: AdamW with betas (0.9, 0.99); weight decay 1e-5 for transformers transformer_trainer.py and 0 for the RVQ vq_trainer.py.
• Learning-rate decay. Paper: not stated. Code: VQ milestones 150000/250000 with gamma 0.05 in the README command vq_option.py README.md; transformer milestone 50000 with gamma 0.1 train_option.py.
• Input width. Code: 263 features per frame for HumanML3D model.py; README says editing needs 263-dimension input README.md. The paper says only D features §3.1.
• Licence of data and models. Code is MIT LICENSE. The README warns that SMPL, SMPL-X, PyTorch3D and the datasets have their own licences README.md. These are not checked here.

Pitfalls

• Paper's own FID numbers disagree. HumanML3D MoMask FID is 0.045 in Table 1 and the introduction, but 0.051 in Table 2 tab. 1 tab. 2 §1. KIT-ML is 0.204 in Table 1 and the introduction, but 0.228 in Table 2 and the abstract tab. 1 tab. 2 §1. Cite the table you use. Multimodal distance differs slightly too: 2.958 vs 2.957 (HumanML3D), 2.779 vs 2.774 (KIT-ML).
• Layer indexing. The method defines V+1 layers numbered 0..V §3.1. The inference figure says the residual transformer predicts t^{2:V} in V-1 steps fig. 3, which uses 1-based indexing. Use 0-based indices in code and check the layer counts against the paper's text.
• Table 2 row labels. The V sweep and the q sweep are not defined in the table itself. The row V = 5 matches the main MoMask row, so V appears to be the number of residual layers on top of the base layer §4.2 tab. 2. Treat this as an inference.
• Step count. The introduction says all layers are generated "within merely 15 iterations" §1. The experiments use L = 10 base iterations §4. The residual step count is V §3.4. The total of 15 is consistent with 10 + 5 but is not stated.
• Guidance scale convention. Use the code's formula with the published scales. Passing the paper's s straight into the code gives a different effective guidance (see Paper vs code).
• Quantization is stochastic during training. The RVQ assignment uses Gumbel noise at temperature 0.5 in code model.py. Reconstructions at eval time are deterministic, so the reported gap between training and evaluation can be larger than the paper implies.
• Commitment weight is not stated. The paper gives no beta eq. 2. The code uses 0.02 with a mean over layers, not the sum in the paper.
• Quantization dropout is per batch in code. A single decision is drawn for every sample in the batch residual_vq.py.
• Masking details matter. Rounding (round vs ceil) and the 80/10/10 split change the effective mask distribution; see Paper vs code.
• Residual sampling is non-uniform in code. A cosine-shaped schedule chooses the training layer transformer.py. A uniform sampler will not reproduce the same training distribution.
• Target length is an input. MoMask needs the motion length, and the authors suggest a separate text-to-length model §5.
• Known weaknesses. The authors report limited diversity, and difficulty with fast root motion such as spinning §5. They also say multimodality should be read together with FID and R-Precision, since a model that outputs random motions can score well on it alone §4.
• Dataset licences. The code is MIT, but HumanML3D, KIT-ML and AMASS have their own terms README.md. Check them before training or distributing derived data.
• Tables are not reproduced from a re-run. The numbers are author-reported. Check the evaluation setup (20 repeats, 95% CI) before comparing with other papers §4.

Lineage

• Builds on: residual quantization from neural audio codecs (SoundStream, AudioLM) and compression §2; generative masked transformers, including MaskGIT, Muse, MAGE and MAGVIT §1 §2; BERT's masked training §3.2; and the VQ-VAE motion tokenizers of TM2T and T2M-GPT, whose codebook update the paper follows §3.1 §2.
• Compared against: TM2T, T2M, MDM, MLD, MotionDiffuse, T2M-GPT, M2DM and ReMoDiffuse tab. 1 tab. 2.
• Superseded by: not stated in this card's sources.