cod-vae-16x4-tiny

COD-VAE 16 x 4 (tiny)

A tiny, decode-first COD-VAE that compresses a 3D shape into 16 latent vectors of 4 dimensions = 64 numbers and decodes them back into an occupancy field. Same latent shape as cod-vae-16x4 and cod-vae-16x4-small, but built for pipelines whose wall clock is the decode forward+backward through a frozen decoder (e.g. reconstruction-reward RL): ~6.5M 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

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.

Held-out reconstruction quality

For reference, cod-vae-16x4-small reaches 0.762 / 0.746 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

import trimesh
from cod_vae import CODVAE

vae = CODVAE.from_pretrained("TimSchneider42/cod-vae-16x4-tiny")

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

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