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

Huaizhi Qu; Mingxin Yang; Xintong Han; Fangneng Zhan; Chunchao Guo

cs.CVOct 10, 2026

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

Authors: Huaizhi Qu, Mingxin Yang, Xintong Han, Fangneng Zhan, Chunchao Guo

Organizations: The Hong Kong University of Science and Technology · Tencent Hunyuan

Abstract

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.

Figures & tables

Figure 1: Improving local responses in VecSet VAEs. (a) Nearby encoder selections replace a shared input. (b) Decoder responses become localized. The models share an input shape and 8192 stored tokens. Lines show most-attended inputs; crosses mark FPS origins. Each slice shows the median token by surface attention, through its volume peak and normalized by that peak.

Panel of Figure 2 — Figure 2: Attention allocation in a standard VecSet VAE. Local selection uses 64 nearest inputs; the dotted line denotes equal allocation. Surface maps (attention ranks 1, 10, median) use each token’s surface peak for log-scale normalization; crosses mark FPS origins. Bars and curves show shape means and 95% intervals; (a, d) share one example.

Panel of Figure 3 — Figure 3: Spatial responses and field orientation in standard VecSet. Shape and token ranks match Figure 2(d) . Slices are normalized by each token’s full-volume peak; crosses mark projected FPS origins. The dashed line gives the trace in (c), alongside shift similarity $C_{\mathrm{rep}}$ . Normal maps saturate at $90^{\circ}$ . Summaries show shape means and 95% intervals.

Panel of Figure 4 — Figure 4: Align-VAE overview. Encoder cross-attention and self-attention produce the $M$ -token code. The decoder prepends $K$ shared learnable tokens before self-attention and retrieves shape tokens with a distance bias. Circled regions show the decoder bank, distance bias, normal-free queries, and normal supervision, respectively. The field receives surface and normal-direction constraints; color bands depict the spatial prior.

Figure 5: VecSet (gray) and Align-VAE (teal): 8192 shape latents, 64 matched shapes. Bars show means and 95% shape-bootstrap intervals. Panel (a) conditions mass on shape tokens.

Method	Tokens	Objaverse			ABO			In-the-wild
		CD $\downarrow$	F1@	F1@	CD $\downarrow$	F1@	F1@	CD $\downarrow$	F1@	F1@
		CD $\downarrow$	0.01 $\uparrow$	0.001 $\uparrow$	CD $\downarrow$	0.01 $\uparrow$	0.001 $\uparrow$	CD $\downarrow$	0.01 $\uparrow$	0.001 $\uparrow$
Dora	4096	5.208	98.90	21.15	5.052	98.49	22.31	4.317	99.97	24.03
HY3D-2.1	4096	6.593	97.72	24.56	6.295	98.28	25.01	11.947	94.92	25.52
VecSet VAE †	4096	4.797	98.97	28.93	5.143	98.74	26.32	3.017	99.83	44.33
VecSet VAE †	8192	3.773	99.62	34.57	3.927	99.54	34.53	2.479	99.95	52.91

Table 1: Median reconstruction scores on Objaverse, ABO, and in-the-wild shapes. Lower CD (

\times 10^{3}

) and higher F1 (%) are better. Throughout the main text, bold and underlined scores mark the best and second-best per column.

\dagger

marks our reproductions of VecSet and LATTICE.

Method	Tokens	Level 1		Level 2		Level 3		Level 4		Overall
Method	Tokens	CD $\downarrow$	F1 $\uparrow$	CD $\downarrow$	F1 $\uparrow$	CD $\downarrow$	F1 $\uparrow$	CD $\downarrow$	F1 $\uparrow$	CD $\downarrow$	F1 $\uparrow$
Dora	4096	3.927	13.57	4.089	21.41	3.932	30.19	4.083	26.25	4.008	22.85
HY3D-2.1	4096	99.117	11.78	66.173	17.04	58.317	15.15	13.855	24.43	59.365	17.10
VecSet VAE †	4096	2.683	69.97	2.132	62.24	3.030	49.40	5.761	33.77	3.402	53.84
VecSet VAE †	8192	2.547	74.21	2.053	64.00	2.732	52.65	11.883	36.61	4.804	56.87
LATTICE †	4096	2.990	64.18	2.937	55.14	3.557	46.29	4.792	29.23	3.569	48.71

Table 2: Mean reconstruction scores on Dora Bench. CD is scaled by

10^{3}

; F1 denotes [email protected] in percent. Overall averages valid objects across all levels. Method notation follows Table 1 .

Figure 6: Local geometry at 8192 tokens on selected shapes. Ground-truth overviews mark (a) rear railings, (b) screw threads, and (c) tire tread. Each row uses the same camera, material, lighting, and crop across methods. Details are rendered directly at higher resolution.

Variant	Component				Metric
	Bank	Distance	Normal-free	Normal	CD $\downarrow$	[email protected] $\uparrow$	S-F $\uparrow$	$E_{\mathrm{angle}}$ $\downarrow$
	Bank	bias	queries	supervision	CD $\downarrow$	[email protected] $\uparrow$	S-F $\uparrow$	$E_{\mathrm{angle}}$ $\downarrow$
A0	$\times$	$\times$	$\times$	$\times$	3.723	35.15	97.30	12.83
A1	$\checkmark$	$\times$	$\times$	$\times$	2.960	42.70	98.30	11.92
A2	$\checkmark$	$\checkmark$	$\times$	$\times$	2.809	45.99	98.44	11.61
A3	$\checkmark$	$\checkmark$	$\checkmark$	$\times$	2.723	46.71	99.11	10.82

Table 3: Cumulative ablation at 4096 tokens; A0–A4 are defined in Section 6.2 . CD is scaled by

10^{3}

; F1 and sharp-point coverage S-F are percentages; normal error

E_{\mathrm{angle}}

is in degrees.

Variant	Local selection	Local mass	Bank mass	Response corr.	Normal error
Variant	$L_{\mathrm{enc}}$ (%) $\uparrow$	$P_{\mathrm{enc},k}$ (%) $\uparrow$	$G_{\mathrm{bank}}$ (%)	$C_{\mathrm{rep}}$ $\downarrow$	$E_{\mathrm{angle}}$ (deg) $\downarrow$
A0	0.076	3.05	0.00	0.392	12.83
A1	0.063	13.80	31.19	0.344	11.92
A2	0.043	9.77	54.27	0.052	11.61
A3	100.000	53.82	53.17	0.034	10.82
A4	100.000	46.73	48.69	0.031	14.48

Table 4: Geometric behavior of A0–A4, defined in Section 6.2 . Local selection measures how often an encoder query selects one of its 64 nearest inputs; local mass sums attention to those neighbors. Lower response correlation and normal error indicate less spatial repetition and better orientation. Bank mass describes attention to shared tokens and has no preferred direction.

Appendix figures & tables13 assets

Supplementary material from the paper’s appendix.

Appendix

Parameter	Meaning	Setting	Experiments
$K$	Decoder bank size	256	A1–A4, Full
$\sigma_{\mathrm{ref}}$	Reference bandwidth	0.097	A2–A4, Full
$M_{\mathrm{ref}}$	Reference token count	4096; 8192	A2–A4; Full
$\beta$	Final prior strength	1	A2–A4, Full
Prior schedule	Delay, then ramp	10k, then 5k	A2–A4, Full
$\lambda_{\mathrm{anchor}}$	Anchor loss weight	0.1	A2–A4, Full

Appendix

Table 5: Training settings for the in-the-wild experiments. Offsets and bandwidths use normalized coordinates; schedule lengths are optimizer steps. Settings are scoped to the listed experiments.

Figure 7: A visual guide to the profiling measurements. (a) Follow each encoder query’s strongest attention link, then count local or repeated selections. (b) Average decoder attention over queries, rank tokens, and accumulate their shares. (c) Hold a token fixed and view its response over spatial locations. (d) Shift a response curve and compare the overlapping values; retain the largest correlation over the specified shifts. (e) Change query normals while fixing the other encoder inputs. (f) Estimate the field gradient with six offset queries and compare its direction with the reference normal. Panels (a,e,f) are schematics. Panels (b–d) reuse the first example in Figures 2 and 3 ; (c,d) use its surface-mass rank-10 token. Panel (d) reports the score for this single trace.

Metric	Symbol	Meaning	Unit
Local selection	$L_{\mathrm{enc}}$	Queries selecting one of 64 neighbors	Fraction
Local mass	$P_{\mathrm{enc},k}$	Attention on 64 nearest inputs	Fraction
Most-selected input	—	Largest share of queries selecting one input	Fraction
Token mass	$p_{i}$	Conditional share of a shape token	Fraction
Top-1% share	$T_{1\%}$	Mass on the top 1% of shape tokens	Fraction
Mass entropy	$H_{\mathrm{shape}}$	Evenness of shape-token use	Unitless

Appendix

Table 6: Lookup for diagnostic names and symbols used in the figures and tables. Coordinate units refer to the normalized geometry. Fractions are displayed as percentages where indicated.

Statistic	Volume queries	Input-surface queries
$100T_{1\%}$	19.50 [18.48, 20.61]	9.83 [9.32, 10.39]
$H_{\mathrm{shape}}$	0.875 [0.871, 0.878]	0.960 [0.959, 0.962]
$R_{\mathrm{shape}}\!\downarrow$	0.953 [0.937, 0.969]	0.304 [0.288, 0.319]
$R_{\mathrm{top}}\!\downarrow$	1.030 [1.001, 1.061]	0.393 [0.370, 0.418]
Uniform reference for $R_{\mathrm{top}}$	1.244 [1.224, 1.266]	0.764 [0.726, 0.804]
$S_{\mathrm{top}}\!\downarrow$ (mass weighted)	0.893 [0.882, 0.905]	0.354 [0.338, 0.371]

Appendix

Table 7: Baseline attention diagnostics: means [95% shape-bootstrap intervals]. Distances and spreads use canonical coordinates; lower values indicate greater spatial locality.

Surface mass rank	Lines through FPS origin $\downarrow$	Lines through volume peak $\downarrow$
1	0.493 [0.456, 0.531]	0.475 [0.434, 0.515]
10	0.548 [0.524, 0.575]	0.518 [0.491, 0.547]
4097 (median)	0.292 [0.268, 0.315]	0.274 [0.250, 0.296]

Appendix

Table 8: Spatial-shift response correlation in the baseline, with 95% confidence intervals. Tokens are ranked by surface attention mass; lines pass through either the FPS origin or the strongest response. Lower correlation indicates less spatial repetition; bold and underlined estimates mark the lowest and second-lowest per column.

Model	Tokens	CD $\downarrow$	[email protected] $\uparrow$	S-CD $\downarrow$	[email protected] $\uparrow$
A0: VecSet	4096	3.723	35.15	3.049	97.30
A1: + LB	4096	2.960	42.70	2.803	98.30
A2: + DB	4096	2.809	45.99	2.759	98.44
A3: + PQ	4096	2.723	46.71	2.480	99.11
A4: + NS	4096	2.526	49.12	2.281	99.50

Appendix

Table 9: Sharp-structure reconstruction for A0–A4 at 4096 tokens. CD and sharp-point distance S-CD are scaled by

10^{3}

; [email protected] and sharp-point coverage [email protected] are percentages.

Model	Tokens	KiB	Mem. $\downarrow$	Encode $\downarrow$	Decode $\downarrow$	KV $\downarrow$	Query $\downarrow$
A0: VecSet	4096	512	3.754	1273.9	53.96	0.466	2.189
A1: + LB	4096	512	3.755	1276.5	56.33	0.484	2.254
A2: + DB	4096	512	3.755	1275.3	56.34	0.495	2.146
A3: + PQ	4096	512	3.755	1275.6	56.46	0.494	2.148
A4: + NS	4096	512	3.755	1275.5	56.49	0.497	2.175
VecSet VAE †	4096	512	2.319	649.7	53.99	0.465	3.806

Appendix

Table 10: Representation storage and inference cost. KiB is stored latent size and Mem. is peak memory in GiB. Encode constructs the latent code; Decode refines it; KV prepares decoder keys and values; Query evaluates 4096 spatial locations. Times are in milliseconds and exclude mesh extraction. Blocks separate cumulative variants and full-scale models at each token budget; memory and times are ranked within each block.

Comparison	$\Delta$ CD $\downarrow$	$\Delta$ [email protected] $\uparrow$	$\Delta E_{\mathrm{angle}}$ $\downarrow$
A0 $\to$ A1	-0.763 [-0.958, -0.455]	7.55 [4.83, 9.47]	-0.91 [-1.02, -0.81]
A1 $\to$ A2	-0.150 [-0.408, 0.055]	3.29 [0.02, 5.76]	-0.31 [-0.35, -0.27]
A2 $\to$ A3	-0.087 [-0.331, 0.107]	0.72 [-2.11, 3.53]	-0.80 [-0.89, -0.71]
A3 $\to$ A4	-0.196 [-0.409, -0.007]	2.42 [0.24, 5.63]	3.66 [3.26, 4.07]

Appendix

Table 11: Successive variant minus preceding variant, with 95% confidence intervals across shapes. CD is scaled by

10^{3}

, F1 changes are percentage points, and normal-error changes are degrees. Negative CD changes and positive F1 changes indicate better reconstruction. Emphasis ranks point estimates of the changes, not statistical significance.

Model	Probe	$T_{1\%}$	$H_{\mathrm{shape}}$	$G_{\mathrm{bank}}$	$R_{\mathrm{shape}}\!\downarrow$	$R_{\mathrm{top}}\!\downarrow$	$S_{\mathrm{sel}}\!\downarrow$
A0	Surface	1.95	0.995	0.00	0.205	0.217	0.343
A0	Volume	5.00	0.985	0.00	0.874	1.197	0.921
A1	Surface	2.66	0.992	31.19	0.175	0.320	0.450
A1	Volume	3.18	0.986	38.43	0.834	0.900	0.926
A2	Surface	2.46	0.990	54.27	0.035	0.028	0.041
A2	Volume	2.00	0.996	97.04	0.079	0.085	0.098

Appendix

Table 12: Decoder attention profiles.

T_{1\%}

is the attention share of the most-read 1% of shape tokens;

H_{\mathrm{shape}}

measures how evenly shape tokens are read.

G_{\mathrm{bank}}

is the shared bank’s attention share.

R_{\mathrm{shape}}

and

R_{\mathrm{top}}

are read distances for all shape tokens and the most-read 1%, respectively;

S_{\mathrm{sel}}

measures spatial response spread. Shares are percentages; distances use normalized coordinates. Distance and spread are ranked separately for surface and volume probes. Mass shares and entropy are descriptive, with no preferred direction.

Model	Query normals	Local selection $\uparrow$	Local mass $\uparrow$	CD $\downarrow$
A0	Clean	0.076	3.05	4.037
A0	Zeroed	0.052	3.04	16.414
A0	Perturbed	0.078	3.05	4.382
A0	Random flips	0.076	2.98	5.759
A2	Clean	0.043	9.77	3.026
A2	Zeroed	0.043	9.87	3.747

Appendix

Table 13: Effect of changing only encoder query normals. A0 is VecSet; A2 includes the bank and distance bias; A3 additionally removes query normals. Local selection

L_{\mathrm{enc}}

and local attention mass

P_{\mathrm{enc},k}

are percentages; CD is scaled by

10^{3}

. Positions and support features remain fixed. Scores are ranked across models under the same query-normal condition.

Model	Angular error $\downarrow$	High-curvature error $\downarrow$	CD $\downarrow$	[email protected] $\uparrow$
A0: VecSet	12.83	21.01	3.723	97.30
A1: + LB	11.92	19.09	2.960	98.30
A2: + DB	11.61	18.50	2.809	98.44
A3: + PQ	10.82	16.84	2.723	99.11
A4: + NS	14.48	18.14	2.526	99.50

Appendix

Table 14: Reconstruction and field orientation. Angular errors are in degrees; the high-curvature subset contains points with the largest local normal variation. CD is scaled by

10^{3}

and sharp-point coverage S-F is a percentage. Lower angular error indicates closer agreement between predicted gradients and reference normals.

Figure 8: Additional qualitative results on selected shapes at 8192 tokens: (a) hull windows, (b) waist canisters and armor plates, (c) vents and rim struts, and (d) handrails and fixtures. Ground-truth overviews locate each detail. All methods share the camera, material, lighting, and crop within each row, with details rendered directly at higher resolution.

Figure 9: Additional qualitative results on selected shapes at 8192 tokens (continued): (e) trouser grooves, (f) wheel spokes and rim, and (g) individual toes. The same rendering protocol is used as in Figure 8 , preserving the size of the enlarged details and the matched views across methods.

Variational Autoencoders

Variational autoencoders are probabilistic encoder–decoder models for learning latent representations, generating data, and estimating uncertainty.

See all 209 papers

MSVS-VAE: Multi-Scale Anchored VecSet for High-Fidelity 3D Reconstruction
Jul 27, 2026Dehao Hao, Kaiyi Zhang, Tanghui Jia +10Variational Autoencoders 3D Reconstruction
The Hong Kong University of Science and Technology, Hong Kong SAR, China · Peking University, Beijing, China · The Hong Kong University of Science and Technology (Guangzhou), Guangzhou, China +1
Beyond Global Latents: Chunk-Based Sparse Grid VAE for Scalable 3D Modeling
Aug 3, 2026Kaiyi Zhang, Zhihao Liang, Haolin Liu +8Variational Autoencoders Autoencoders
Hong Kong University of Science and Technology · Tencent Hunyuan · The Chinese University of Hong Kong
TopoMesh: High-Fidelity Mesh Autoencoding via Topological Unification
Mar 25, 2026Guan Luo, Xiu Li, Rui Chen +6Autoencoders Variational Autoencoders
Tsinghua University · ByteDance Seed · HKUST

Jul 27, 2026J/K move · Enter open · S save

MSVS-VAE: Multi-Scale Anchored VecSet for High-Fidelity 3D Reconstruction

Dehao Hao, Kaiyi Zhang, Tanghui Jia, Xiangjun Gao+9

The Hong Kong University of Science and Technology, Hong Kong SAR, China · Peking University, Beijing, China · The Hong Kong University of Science and Technology (Guangzhou), Guangzhou, China+1

Variational Autoencoders 3D Reconstruction Generative Modeling 3D Representation Learning

Open paper AskPDF arXiv

High-fidelity 3D generative modeling increasingly relies on the latent diffusion paradigm, where the reconstruction quality of the underlying 3D VAE becomes a primary bottleneck. Existing approaches largely follow two paradigms: sparse voxel-based representations achieve strong reconstruction quality but incur significant memory and computational overhead, while set-based representations are compact and continuous yet typically lag in fidelity due to latent sparsity and excessive global smoothness. We propose MSVS-VAE, a hierarchical set-based VAE that closes this fidelity gap without sacrificing compactness. Our key idea is to progressively densify anchored VecSet latents via hierarchical point-shuffle upsampling, increasing spatial capacity for fine-grained geometry modeling. To efficiently decode from the densified hierarchy, we replace global cross-attention with AVS-Conv, a geometry-aware local aggregation operator operating within local neighborhoods rather than the exhaustive latent set. We further introduce multi-scale query decoding to fuse coarse-to-fine latent features, where coarse scales provide stable global context, and fine scales refine localized geometry, reducing artifacts from overly local receptive fields. Extensive experiments on Objaverse, ABO, and in-the-wild benchmarks demonstrate that MSVS-VAE consistently outperforms prior set-based and voxel-based VAEs, delivering approximately 10x faster decoding than prior set-based methods and approximately 10x higher compactness than voxel-based baselines.

Abstract

Variational Autoencoders

Explore similar work

MSVS-VAE: Multi-Scale Anchored VecSet for High-Fidelity 3D Reconstruction

Beyond Global Latents: Chunk-Based Sparse Grid VAE for Scalable 3D Modeling

TopoMesh: High-Fidelity Mesh Autoencoding via Topological Unification