MissawB/creature-flow-point-dit

Creature Flow: skeleton-conditioned Point-DiT

A flow matching model that generates a 2048-point surface of a quadruped from its kinematic skeleton alone (35 joints, SMAL-like topology). Because it only reads the skeleton, it can also grow bodies for skeletons it was never trained on, such as a chimera grafted from a lion body and a horse neck.

Try it: creature-flow-viewer plays the whole flow from Gaussian noise (t = 0) to the finished specimen (t = 1).

2048-point clouds generated in 10 Euler steps from the skeleton alone (red). The last column is a chimera: lion body, horse neck and head.

Model

Architecture	Point-DiT, 6 blocks, width 256, 8 heads, AdaLN-Zero time modulation
Conditioning	Transformer skeleton encoder, then point-to-joint cross-attention with a multi-scale distance bias
Parameters	10.4 M
Objective	Flow matching on linear paths, $x_t = (1-t),x_0 + t,x_1$, target $x_1 - x_0$, with per-shape OT coupling: the 2048 noise points are matched to the 2048 surface points by exact assignment
Sampling	Euler ODE solver, 10 steps
Weights	Exponential moving average (decay 0.995) after 180 epochs
Data scale	Points and joints are multiplied by 4 inside the flow (data_scale), so the data matches the unit-variance noise
Training data	procedural-quadrupeds, 500 shapes
Training cost	About 32 min on an RTX 4060 Laptop GPU in bf16, plus about 30 min of CPU once to precompute the OT-matched noise

Each attention head $h$ subtracts $\lambda_h \lVert \hat{x}_{1,i} - J_k \rVert^2$ from the point-to-joint score, with $\lambda_h = 2 \cdot 2^h$, so the 8 heads range from near-global (2) to very local (256) attention. The distance uses an estimate of the final position $\hat{x}_1 = x_t + (1-t),\hat{v}$ rather than $x_t$, which is pure noise early in the trajectory.

Results

Evaluated on 40 skeletons never seen in training (8 per family). Each generated shape is assigned to the family whose 8 reference shapes are closest on average in Chamfer distance.

	This checkpoint
Family recognised	40/40 (100 %)
Chamfer distance to its family	0.0009
MMD-CD (lower is better)	0.00084
COV-CD (higher is better)	57 %
1-NNA-CD (50 % is ideal)	51.0 % (generated 2 %, real 100 %)
Median distance to the true surface	0.0029 (real shapes: 0.0004)

For scale, two real shapes of the same family are 0.0007 apart on average, and an untrained model scores about 1.9. Two training runs of this recipe both recognised 40/40, with a Chamfer distance of 0.0009 and a median surface distance of 0.0029 and 0.0030. MMD, COV and 1-NNA are the standard point-cloud generation metrics of PointFlow, computed between 200 generated shapes and 200 real shapes drawn separately (scripts/evaluate_generation.py). The surface distance is the signed distance of each point to the nearest bone cylinder, computed from the skeleton.

In the ablation, the multi-scale distance bias accounts for most of the gain: a single fixed bias reached only 40 %.

The last step came from the data, not the model. The model only reads the skeleton, and bovids differed from horses mainly by body thickness, which no joint encoded: about half of the generated bovids were taken for horses (90 % overall). Generator version 4 makes that difference visible in the skeleton. Legs stand wider under a thicker body, and bovids carry the neck at 20° instead of 49°. With the same model and training, bovids went from 13/24 to 24/24 over three seeds, with no loss on the other families.

The surfaces were then too thick. Normalised shapes have a standard deviation of about 0.25, while the starting noise has a standard deviation of 1. Multiplying the coordinates by 4 inside the flow (data_scale: 4, in the spirit of EDM's $\sigma_{data}$) improves every measure over three seeds each:

Mean over 3 seeds	data_scale: 1	data_scale: 4
Median distance to the true surface	0.0129	0.0092
Points more than 30 % of the radius outside the surface	15.9 %	8.8 %
MMD-CD	0.00224	0.00154
COV-CD	48 %	51 %
1-NNA-CD	98.4 %	96.3 %

A scale of 8 (one run) was worse than 4.

The remaining blur came from undertraining. More ODE steps changed nothing, and Fourier input features were slightly worse, but shapes kept improving with every training stage while the loss barely moved (1.30 to 1.29). Training 90 epochs instead of 30 fixed most of it, over three runs each:

Mean over 3 seeds (data_scale: 4)	30 epochs	90 epochs
Median distance to the true surface	0.0090	0.0051
Points more than 30 % of the radius outside the surface	8.8 %	2.2 %
MMD-CD	0.00150	0.00098
COV-CD	50 %	56 %
1-NNA-CD	99.6 %	53.6 %
of which generated shapes correctly classified	99 %	7 %
of which real shapes correctly classified	100 %	100 %

A 1-NNA near 50 % does not mean the generated shapes are indistinguishable here. It averages two opposite errors. A real shape always has another real shape as its nearest neighbour (100 %), because the more diffuse generated shapes stay farther from it than its real neighbours. A generated shape, in turn, almost always has a real shape as its nearest neighbour (7 %): it sits closer to the real shapes than to other generated ones. The gap that remains is visible in the split, which evaluate_generation.py prints.

Per-shape OT coupling

After 90 epochs the surface plateaued at a median distance of about 0.005. A cosine learning-rate schedule, a wider model (width 384, 23 M parameters), scaling t by 1000 before the sinusoidal embedding, and a k-nearest-neighbour local attention in every block all stayed between 0.0049 and 0.0054. The gap was the same on training and unseen skeletons, and the same width (0.005 to 0.008) for bone radii from 0.04 to 0.32, which rules out memorisation and scale ambiguity.

What was left was the training target. With independent coupling, each noise point is paired with a random surface point, so $x_1 - x_0$ has a huge irreducible variance and the surface detail is a tiny part of it. Matching noise to points by optimal transport inside each shape makes the target nearly deterministic:

90 epochs, bf16	Independent coupling	OT + 20 % independent noise	Pure OT (2 runs)
Final loss	1.25	0.29	0.013
Median distance to the true surface	0.0052	0.0094	0.0040 / 0.0042
Mean signed distance (outward bias)	+0.0019	+0.0033	+0.0003 / +0.0002
MMD-CD	0.00097	0.00162	0.00092 / 0.00095
Family recognised	100 %	100 %	100 % / 100 %

Any share of independent noise brings the variance back. Paths also become nearly straight: with this checkpoint, Euler with 5 steps (5 network evaluations, surface distance 0.0039) matches Heun with 20 steps (40 evaluations, 0.0040), and Euler 10 is the default. Heun 2 still gives usable shapes (MMD 0.00124, surface distance 0.0071). These step counts were measured on a 90-epoch checkpoint.

With a nearly deterministic target, the loss keeps falling well after 90 epochs, and the surface sharpens with it (200 vs 200 shapes, Euler 10):

	90 epochs	180 epochs (2 runs)	360 epochs
Median distance to the true surface	0.0038	0.0029 / 0.0030	0.0026
MMD-CD	0.00092	0.00084 / 0.00085	0.00082
COV-CD	58 %	57 % / 56 %	56 %

A constant learning rate does as well as a cosine decay (0.0030). The gain keeps going but slows down after 180 epochs, which is the default. Two OT noise draws per shape instead of eight give the same result.

Use

git clone https://huggingface.co/MissawB/creature-flow-point-dit
cd creature-flow-point-dit
pip install -r requirements.txt safetensors

import torch
from safetensors.torch import load_file
from data.smal_quadrupeds import SMALQuadrupedGenerator
from models.point_dit import PointDiT
from models.flow_matching import OptimalTransportFlowMatching

model = PointDiT(point_dim=3, hidden_dim=256, num_layers=6, num_heads=8,
                 num_joints=35, spatial_distance_scale=2.0 / 4**2, data_scale=4.0)
model.load_state_dict(load_file("model.safetensors"))
model.eval()

joints = SMALQuadrupedGenerator(seed=0).generate_instance("felidae")["joints"]
fm = OptimalTransportFlowMatching(model)
_, points = fm.sample(torch.from_numpy(joints).unsqueeze(0), num_points=2048,
                      num_steps=10, solver="euler")
print(points.shape)  # (1, 2048, 3)

To build a chimera, graft a donor subtree onto a host skeleton with pipeline.hybrid_splicer.KinematicHybridSplicer().splice(host, donor, junction_joint="neck_base") and sample from the result.

The repository also holds the training (scripts/train.py), evaluation (scripts/evaluate_families.py, scripts/evaluate_generation.py) and viewer export (scripts/sample_and_export.py) scripts. Those scripts read .pt checkpoints produced by training. This repo ships only the EMA weights as model.safetensors.

Limitations

• Synthetic data. Shapes are cylinders sampled around a parametric SMAL-like skeleton, not real animal scans or the SMAL model itself. Within-family variation is small (about ±3 % on proportions).
• An easy test. The five families are separated by construction (proportions, stance, head carriage) and vary little inside a family, so 100 % shows that the model follows the skeleton, not that it generalises to real animals.
• Surface precision. Generated surfaces are still more diffuse than real ones (median distance 0.0029 against 0.0004), and shapes about 1.3 times as far from their family as two real shapes are from each other (0.0009 against 0.0007). A nearest-neighbour classifier still recognises every real shape (1-NNA on the real side: 100 %).
• Grafting aligns the donor subtree by translation only, with no scaling or rotation.

Collection

Source code, experiments and full results: GitHub.

Part of the Skeleton-conditioned creature flow collection with the viewer Space and the training dataset.