README.md

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

# COD-VAE 4 x 4 (tiny)

A tiny, decode-first [COD-VAE](https://arxiv.org/abs/2503.08737) 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](https://huggingface.co/TimSchneider42/cod-vae-4x4) and
[cod-vae-4x4-small](https://huggingface.co/TimSchneider42/cod-vae-4x4-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`](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. The JAX decode numbers below
include the channel-last plane layout (cod-vae >= 56b2c82).

## Architecture vs the -small recipe

| | cod-vae-4x4-small | this model |
|---|---|---|
| embed dim / heads | 256 / 4 | 128 / 4 |
| encoder | 3 blocks x 3 layers, 512 patches, mlp 4 | 2 blocks x 2 layers, 256 patches, mlp 2 |
| refinement decoder | 6 layers, 16-px patches | 4 layers, 32-px patches |
| query planes (`query_dim`) | 16 channels at 128² | 8 channels at 96² |
| latent decoder layers | 12 | 6 |
| total parameters | ~35M | ~6.5M |

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.

| model | step | throughput |
|---|---|---|
| cod-vae-16x8 (full) | ~350 ms | 2.9k shapes/s |
| cod-vae-16x8-small | 43.5 ms | 23.6k shapes/s |
| **cod-vae-16x8-tiny** | **8.0 ms** | **127k shapes/s** |

## Held-out reconstruction quality

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

For reference, cod-vae-4x4-small reaches 0.650 / 0.698 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

```python
import trimesh
from cod_vae import CODVAE

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

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:

```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 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](https://github.com/TimSchneider42/cod-vae/blob/main/TRAINING.md)
for the exact commands.

## Citation

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