粗格错了细格来纠:JoVox 联合流匹配生成体素,训练吞吐 8 倍于 TRELLIS.2

JoVox: Joint Hierarchical Flow Matching for Lightweight Voxel Generation

Xiangjun Tang, Jianqi Chen, Peter Wonka

cs.CV, cs.GT

2026-11-07

KAUST 提出 JoVox,用离散流匹配联合建模体素从 4³ 到 64³ 的所有层级,让细格参与粗格决策;比 TRELLIS.2 训练吞吐高 8 倍、体素 Chamfer 低 11%。

这篇在解决什么

体素(voxel)记录 3D 形状占据空间中哪些格子,是 3D 生成里最轻量的中间表示之一,TRELLIS 系列就是先生成 64³ 体素再细化成 mesh。生成高保真体素的两条现有路线各有硬伤。

潜空间路线(TRELLIS、TRELLIS.2)先用 autoencoder 把体素压成紧凑 latent 网格再生成,保真度高,但要多训一个模型,且生成器仍要处理 latent 网格的全部格子,包括空白区域,开销随分辨率三次方增长。层级路线(XCube、OctFusion、OctGPT 等)从 4³ 粗网格逐级细化到 64³,只处理被占据的格子,轻量,但有个结构性缺陷:决定所有细格能在哪出现的粗层,恰恰在对形状了解最少的时候生成,而且定下后不再修改。论文的评语是「plausible but wrong」,看似合理实则错位,细层只能继承这些错误。

方法

JoVox 的核心改动是学整个层级的联合分布,替代逐级各自建模,让细格参与粗格的决策。落到三个组件:

网络与 TRELLIS.2 的体素阶段同架构,transformer 把所有层级 token 放进一个序列做 full self-attention,图像条件经 DINOv3 特征做 cross-attention。训练时各层级的遮蔽比例独立采样,网络见过任意比例组合,推理中的任何中间状态都落在训练分布内;loss 先在每层级内平均、再跨层级平均,防止格子数多得多的细层主导,小遮蔽比例(预测最不确定的区域)采样得更频繁。

结果

图像条件生成:训练用 TRELLIS-500K 的 174,528 个资产,测试在 1,000 个训练未见过的 Toys4k 物体上。

方法CDV ↓IoU ↑LPIPSV ↓CDshell(×10⁻³)↓
TRELLIS.20.012500.40170.15240.524
JoVox0.011070.41580.13950.407

体素 Chamfer 低 11%,经同一个共享 mesh 阶段细化后外表面 Chamfer 低 22%,从体素到 mesh 的全部八项指标均排第一。无条件生成在 ShapeNet 四类(chair、airplane、table、rifle)平均后六项指标全优:FIDV 2.06 对最好基线 10.1,低 80%;FIDVD 15.3,低 64%;1-NNA 0.574,最接近理想值 0.5。

效率差距更大。同 backbone、单张 A100 80GB、batch 32:峰值显存 29.0 GB 对 47.3 GB,训练吞吐 8 倍;单次前向平均 0.82 TFLOPs 对 12.2 TFLOPs,约 1/15。原因是训练状态平均只有 273 个格子带 token,TRELLIS.2 却要为每个形状处理 16³ latent 网格的全部 4096 格。完整训练用 4 张 H100:约两天、80K 步,在全部指标上超过训练 400K 步的 TRELLIS;约两天半追平用 32 张 H100 训练的 TRELLIS.2。

消融把贡献拆开:单层预测 CDV 0.02150,换两层预测降到 0.01941,再加联合建模降到 0.01506,联合建模这一步贡献最大。采样策略对比 MDLM、ReMDM、ReMDM-conf、MaskGIT,重遮蔽方案的 CDV 与 IoU 均最优。

为什么重要

对做 3D 生成的人,这是一条质量压过 SOTA 潜空间方法、成本低一个量级的路线,4 张卡能复现完整训练,学术组和小团队够得着。对更广的生成建模读者,「早期决策允许被后续信息修正」不限于体素,任何层级式生成都面临粗层锁定问题,联合分布加重遮蔽是可迁移的解法。也要直说:11% 的 Chamfer 提升属渐进,这篇的分量主要在效率侧。

局限与存疑

作者自述两条:训练集约 175K 资产,不到 TRELLIS.2(约 800K)的四分之一,更多数据应仍有提升空间;同一步内多个 occupancy code 独立采样仍可能不一致,重遮蔽只是缓解,没有根除。另有两点存疑:「两天半追平 TRELLIS.2」依赖对方未公开的训练步数,只能算近似对齐;推理速度对比与 256³ 分辨率结果正文只有附录引用,没有数字,跟进前先翻附录核对。

术语

原文与代码

社区讨论

相关论文

全部论文解读