---
license: mit
library_name: cod-vae
pipeline_tag: feature-extraction
tags:
- 3d
- shape-reconstruction
- autoencoder
- vae
- occupancy
---

# COD-VAE 8 x 16 (small)

A compact, decode-optimized [COD-VAE](https://arxiv.org/abs/2503.08737) that
compresses a 3D shape into **8 latent vectors of 16 dimensions = 128 numbers**
and decodes them back into an occupancy field. Same latent shape as
[cod-vae-8x16](https://huggingface.co/TimSchneider42/cod-vae-8x16), 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-8x16](https://huggingface.co/TimSchneider42/cod-vae-8x16), the two
> models define **different latent spaces** — latents from one cannot be decoded with
> the other.

Trained with [`cod-vae`](https://github.com/TimSchneider42/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-8x16

| | cod-vae-8x16 | 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

```python
import trimesh
from cod_vae import CODVAE

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

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

```python
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](https://huggingface.co/datasets/TimSchneider42/tactile-mnist-abc-dataset-small),
and all 11,480
[tactile-mnist-mnist3d](https://huggingface.co/datasets/TimSchneider42/tactile-mnist-mnist3d)
meshes. Only training splits; meshes preprocessed with the original authors'
[sdf_gen](https://github.com/1zb/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](https://github.com/TimSchneider42/cod-vae/blob/main/TRAINING.md#how-the-published-cod-vae-16xm-small-models-were-trained)
for the exact commands.

## Held-out reconstruction quality

| source | held-out shapes | volume IoU | near-surface accuracy |
|---|---|---|---|
| ABC (CAD parts) | 128 | 0.8367 | 0.8082 |

For reference, the full-size cod-vae-8x16 reaches 0.854 / 0.829 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:

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