File size: 4,100 Bytes
cfaa072 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | ---
license: mit
library_name: cod-vae
pipeline_tag: feature-extraction
tags:
- 3d
- shape-reconstruction
- autoencoder
- vae
- occupancy
---
# COD-VAE 4 x 8 (tiny)
A tiny, decode-first [COD-VAE](https://arxiv.org/abs/2503.08737) that compresses a
3D shape into **4 latent vectors of 8 dimensions = 32 numbers** and
decodes them back into an occupancy field. Same latent shape as
[cod-vae-4x8](https://huggingface.co/TimSchneider42/cod-vae-4x8) and
[cod-vae-4x8-small](https://huggingface.co/TimSchneider42/cod-vae-4x8-small),
but built for pipelines whose wall clock is the decode forward+backward through a
frozen decoder (e.g. reconstruction-reward RL): ~6.6M 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-4x8-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.6M |
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.6671 | 0.7102 |
For reference, cod-vae-4x8-small reaches 0.733 / 0.743 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-4x8-tiny")
mesh = trimesh.load("bunny.obj", force="mesh")
latent, transform = vae.encode_mesh(mesh, return_transform=True) # (4, 8)
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}
}
```
|