TimSchneider42/cod-vae-4x4-small

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

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.

Usage

import trimesh
from cod_vae import CODVAE

vae = CODVAE.from_pretrained("TimSchneider42/cod-vae-4x4-small")

mesh = trimesh.load("bunny.obj", force="mesh")
latent, transform = vae.encode_mesh(mesh, return_transform=True)   # (4, 4)
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:

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

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:

@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}
}