COD-VAE 4 x 4 (small)
A compact, decode-optimized
COD-VAE that
compresses a 3D shape into
4 latent vectors of 4 dimensions = 16 numbers
and decodes them back into an occupancy field. Same latent shape as
cod-vae-4x4, but a ~5x
smaller network tuned for fast decoding — including the backward pass, for pipelines
that train through the frozen decoder: ~39M parameters instead of 188M, with
a ~20M decode path instead of 90M.
Note: although the latent shape matches
cod-vae-4x4, the two
models define
different latent spaces — latents from one cannot be decoded with
the other.
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.
Architecture vs cod-vae-4x4
| cod-vae-4x4 | this model |
|---|
| embed dim / heads | 512 / 8 | 256 / 4 |
| encoder | 4 blocks x 3 layers | 3 blocks x 3 layers |
| refinement decoder | 12 layers, 8-px patches (769 tokens) | 6 layers, 16-px patches (193 tokens) |
| latent decoder layers | 12 | 12 |
query planes (query_dim) | 32 channels | 16 channels |
| total parameters | 188M | ~39M |
| decode-path parameters | 90M | ~20M |
The shipped config also pins attention_implementation="default" (the XLA path):
on the short decode sequences of this architecture it is ~1.3x faster than letting
"auto" pick cuDNN's fused kernel.
Decode speed (H100, JAX float16, measured on the 16x8 variant)
num_latents and latent_dim barely move the decode cost, so these numbers hold
for the whole -small family.
| cod-vae-16x8 | 16x8-small |
|---|
| forward+backward through the full latent, batch 1024 x 2048 queries | ~350 ms (2.9k shapes/s) | 43.5 ms (23.6k shapes/s) |
| end-to-end in a tactile RL training loop (measured, 50k-step arms) | 1.53 env-steps/s | 11.75 env-steps/s |
Usage
1import trimesh
2from cod_vae import CODVAE
3
4vae = CODVAE.from_pretrained("TimSchneider42/cod-vae-4x4-small")
5
6mesh = trimesh.load("bunny.obj", force="mesh")
7latent, transform = vae.encode_mesh(mesh, return_transform=True) # (4, 4)
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
The same merged dataset of 110,077 shapes used for the full-size grid, built with the
cod-vae-dataset tool: the 48,597 ShapeNet training shapes (3DShape2VecSet
preprocessing, 55 synsets), 50,000 CAD meshes from
tactile-mnist-abc-dataset-small,
and all 11,480
tactile-mnist-mnist3d
meshes. Only training splits; meshes preprocessed with the original authors'
sdf_gen recipe.
Training recipe
The architecture was selected in an ablation campaign against a hard quality floor
(held-out ABC IoU >= 0.83 for the 16x8 configuration), then retrained as a grid.
Two stages, both with the reference hyperparameters unless noted:
| stage 1 (autoencoder) | stage 2 (latent VAE) |
|---|
| epochs | 200 (one trunk per num_latents, shared by its row) | 100 |
| batch | 128 per GPU x 2 GPUs = 256 | 256 per GPU x 2 GPUs = 512 |
| learning rate | 1e-4, scaled by effective batch / 256 | same, halved at epochs 60/70/80/90 |
| dataset repeat | 8 per epoch | 8 per epoch |
| precision | float32 with TF32 matmuls | same |
Doubling stage 1 from the reference 100 to 200 epochs was measured worth +0.009 trunk
IoU (~+0.003 after stage 2). See the
training guide
for the exact commands.
Held-out reconstruction quality
| source | held-out shapes | volume IoU | near-surface accuracy |
|---|
| ABC (CAD parts) | 128 | 0.6500 | 0.6979 |
For reference, the full-size cod-vae-4x4 reaches 0.671 / 0.712 on ABC — the ~8x decode
speedup costs 0.02-0.03 IoU. Measured on the ABC test split, which is disjoint from
training, on the decoded occupancy field: IoU over points drawn uniformly from the
cube, accuracy over points drawn near the surface.
Citation
The model architecture and training recipe are from:
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}
