# Align-VAE: Aligning Attention and Geometry in VecSet VAEs for 3D Reconstruction

arXiv: https://arxiv.org/abs/2610.14734 · HTML: https://mlsanity.com/paper/2610.14734/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 17 of its 20 references that are in the map.
- First posted on arXiv 10 Oct 2026.

## Builds on (most cited first, 17 of 17)

- [Efficient Streaming Language Models with Attention Sinks](https://mlsanity.com/paper/2309.17453/) (2023)
- [Objaverse: A Universe of Annotated 3D Objects](https://mlsanity.com/paper/2212.08051/) (2022)
- [Vision Transformers Need Registers](https://mlsanity.com/paper/2309.16588/) (2023)
- [Structured 3D Latents for Scalable and Versatile 3D Generation](https://mlsanity.com/paper/2412.01506/) (2024)
- [3DShape2VecSet: A 3D Shape Representation for Neural Fields and Generative Diffusion Models](https://mlsanity.com/paper/2301.11445/) (2023)
- [CLAY: A Controllable Large-scale Generative Model for Creating High-quality 3D Assets](https://mlsanity.com/paper/2406.13897/) (2024)
- [ABO: Dataset and Benchmarks for Real-World 3D Object Understanding](https://mlsanity.com/paper/2110.06199/) (2021)
- [Michelangelo: Conditional 3D Shape Generation based on Shape-Image-Text Aligned Latent Representation](https://mlsanity.com/paper/2306.17115/) (2023)
- [TripoSG: High-Fidelity 3D Shape Synthesis using Large-Scale Rectified Flow Models](https://mlsanity.com/paper/2502.06608/) (2025)
- [XCube: Large-Scale 3D Generative Modeling using Sparse Voxel Hierarchies](https://mlsanity.com/paper/2312.03806/) (2023)
- [Hunyuan3D 2.1: From Images to High-Fidelity 3D Assets with Production-Ready PBR Material](https://mlsanity.com/paper/2506.15442/) (2025)
- [Direct3D-S2: Gigascale 3D Generation Made Easy with Spatial Sparse Attention](https://mlsanity.com/paper/2505.17412/) (2025)
- [Dora: Sampling and Benchmarking for 3D Shape Variational Auto-Encoders](https://mlsanity.com/paper/2412.17808/) (2024)
- [3DILG: Irregular Latent Grids for 3D Generative Modeling](https://mlsanity.com/paper/2205.13914/) (2022)
- [Sparc3D: Sparse Representation and Construction for High-Resolution 3D Shapes Modeling](https://mlsanity.com/paper/2505.14521/) (2025)
- [LATTICE: Democratize High-Fidelity 3D Generation at Scale](https://mlsanity.com/paper/2512.03052/) (2025)
- [MSVS-VAE: Multi-Scale Anchored VecSet for High-Fidelity 3D Reconstruction](https://mlsanity.com/paper/2607.24436/) (2026)

## Abstract (arXiv)

Set-based variational autoencoders (VecSet VAEs) represent 3D shapes as compact sets of latent tokens. Although each token originates from a query at a sampled surface location, whether these latents truly align with geometry remains unclear. We investigate this alignment and identify four issues: (i) attention-sink-like tokens that attract queries from across the shape; (ii) spatially periodic repetitions in decoder attention responses; (iii) mismatched normal information in encoder and decoder query features; and (iv) missing supervision of the intrinsic relation between signed distance fields and surface normals. To address these issues, we introduce Align-VAE with four corresponding modifications: (i) a shared bank of learnable decoder tokens provides a workspace for global computation; (ii) distance-biased cross-attention guides retrieval toward nearby predicted spatial anchors; (iii) normal-free encoder queries align the two stages' access to normal information while retaining normals in the encoder keys and values; and (iv) a screened Poisson-inspired regularizer supervises surface values and normal-direction derivatives. These changes align latent retrieval and field prediction with geometry without enlarging the per-shape code. On the ABO dataset, Align-VAE reduces Chamfer distance by 26.6% compared to the standard VecSet VAE.