COD-VAE 16 x 8 (tiny)
A tiny, decode-first
COD-VAE that compresses a
3D shape into
16 latent vectors of 8 dimensions = 128 numbers and
decodes them back into an occupancy field. Same latent shape as
cod-vae-16x8 and
cod-vae-16x8-small,
but built for pipelines whose wall clock is the decode forward+backward through a
frozen decoder (e.g. reconstruction-reward RL): ~6.6M parameters, roughly
4x faster than -small and 33x faster than the full-size model.
Note: the latent shape matches its siblings, but every model defines its own
latent space — latents from one cannot be decoded with another.
Trained with
cod-vae, a PyTorch/JAX
reimplementation of COD-VAE (Cho et al., ICCV 2025). The weights are a
self-contained npz and load with either backend. The JAX decode numbers below
include the channel-last plane layout (cod-vae >= 56b2c82).
Architecture vs the -small recipe
| cod-vae-16x8-small | this model |
|---|
| embed dim / heads | 256 / 4 | 128 / 4 |
| encoder | 3 blocks x 3 layers, 512 patches, mlp 4 | 2 blocks x 2 layers, 256 patches, mlp 2 |
| refinement decoder | 6 layers, 16-px patches | 4 layers, 32-px patches |
query planes (query_dim) | 16 channels at 128² | 8 channels at 96² |
| latent decoder layers | 12 | 6 |
| total parameters | ~35M | ~6.6M |
The shipped config pins attention_implementation="default" (the XLA path), which
is measurably faster than cuDNN's fused kernel on these short sequences.
Decode speed (H100, JAX float16, batch 1024 x 2048 queries, fwd+bwd through the full latent)
num_latents and latent_dim barely move the decode cost, so these numbers
(measured on the 16x8 variant) hold for the whole -tiny family.
| model | step | throughput |
|---|
| cod-vae-16x8 (full) | ~350 ms | 2.9k shapes/s |
| cod-vae-16x8-small | 43.5 ms | 23.6k shapes/s |
| cod-vae-16x8-tiny | 8.0 ms | 127k shapes/s |
Held-out reconstruction quality
| source | held-out shapes | volume IoU | near-surface accuracy |
|---|
| ABC (CAD parts) | 128 | 0.7665 | 0.7497 |
For reference, cod-vae-16x8-small reaches 0.842 / 0.804 on the same protocol — the extra
~4x decode speedup costs additional quality. The 16x8 configuration was qualified
against a hard floor of 0.75 ABC volume IoU before the grid was trained.
Usage
1import trimesh
2from cod_vae import CODVAE
3
4vae = CODVAE.from_pretrained("TimSchneider42/cod-vae-16x8-tiny")
5
6mesh = trimesh.load("bunny.obj", force="mesh")
7latent, transform = vae.encode_mesh(mesh, return_transform=True) # (16, 8)
8reconstruction = vae.decode_mesh(latent, transform=transform) # trimesh.Trimesh
Latents can also be computed from raw surface point clouds and decoded at arbitrary
query points:
1latents = vae.encode(points) # (N, 3) in [-1, 1]^3
2logits = vae.decode(latents, queries) # occupancy logits, positive inside
3volume = vae.decode_volume(latents, resolution=128) # dense logit grid
Install with pip install cod-vae[torch,hub] (or cod-vae[jax,hub]).
Training data and recipe
Same 110,077-shape merged dataset and two-stage recipe as the
-small grid: a
200-epoch stage-1 trunk per
num_latents (shared by its row) and a fresh 100-epoch
stage 2 per cell with 6 latent-decoder layers; see the
training guide
for the exact commands.
Citation
1@inproceedings{cho2025cod,
2 author={Cho, In and Yoo, Youngbeom and Jeon, Subin and Kim, Seon Joo},
3 title={Representing 3D Shapes with 64 Latent Vectors for 3D Diffusion Models},
4 booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)},
5 year={2025}
6}
