TimSchneider42/cod-vae-16x16-small

COD-VAE 16 x 16 (small)

A compact, decode-optimized COD-VAE that compresses a 3D shape into 16 latent vectors of 16 dimensions = 256 numbers and decodes them back into an occupancy field. Same latent shape as cod-vae-16x16, 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-16x16, 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-16x16

	cod-vae-16x16	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)

latent_dim changes only the tiny latent projections, so these numbers hold for the whole -small family.

	cod-vae-16x8	this model
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

import trimesh
from cod_vae import CODVAE

vae = CODVAE.from_pretrained("TimSchneider42/cod-vae-16x16-small")

mesh = trimesh.load("bunny.obj", force="mesh")
latent, transform = vae.encode_mesh(mesh, return_transform=True)   # (16, 16)
reconstruction = vae.decode_mesh(latent, transform=transform)      # trimesh.Trimesh

Latents can also be computed from raw surface point clouds and decoded at arbitrary query points:

latents = vae.encode(points)                          # (N, 3) in [-1, 1]^3
logits = vae.decode(latents, queries)                 # occupancy logits, positive inside
volume = 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 (the 16-latent trunk is shared by the whole 16x* 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.8721	0.8302

For reference, the full-size cod-vae-16x16 reaches 0.903 / 0.863 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:

@inproceedings{cho2025cod,
  author={Cho, In and Yoo, Youngbeom and Jeon, Subin and Kim, Seon Joo},
  title={Representing 3D Shapes with 64 Latent Vectors for 3D Diffusion Models},
  booktitle={Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)},
  year={2025}
}