跳到主要内容
返回时间线
arXiv来源发表:

体素生成不再逐层定形:JoVox让细层参与粗层决策,训练显存与吞吐同时改善

核心概要

JoVox把体素八叉树的各层放进同一个离散流匹配框架联合建模,让已经生成的细层单元参与尚未确定的粗层单元预测,而不是像分层方法那样先定粗层、再不回头地逐层细化。它同时一次预测细两级的单元,把网络要处理的层数减半。在Toys4k图像条件生成上,其体素倒角距离为0.01107、IoU为0.4158,优于TRELLIS.2的0.01250与0.4017;在相同骨干下,训练所需显存更低、吞吐更高。

AI-generated editorial illustration: JoVox: Joint Hierarchical Flow Matching for Lightweight Voxel Generation

深度剖析

体素生成现在可以让细层单元参与粗层单元的判定,粗层不再是一次定下、之后无法修正的骨架。 此前的分层方法先生成粗网格,再逐层细化被占据的单元,粗层只依据更粗的层决定,且从不回访,因此常常看似合理却是错的,细层会继承这些错误。 在Toys4k图像条件生成上,联合建模的完整模型体素倒角距离为0.01506、IoU为0.334,而同样采用两级预测但逐层生成层级的变体为0.01941与0.264,两者共享网络结构和训练数据、训练步数相同。

生成过程从全掩码状态出发,每一步在任意层级揭开少量单元,层级随着单元被揭开而生长。 掩码离散流匹配原本用于长度固定的序列,而体素层级会随单元揭开而扩张,这里把该方法扩展到这种会生长的层级上。 模型是一个transformer,为每个单元取一个token,把各层token放进同一序列做全自注意力;训练时各层的掩码比例独立抽取,因此任何揭序策略产生的中间状态都落在训练分布内。

一次预测细两级的单元,使网络只需处理一半的层级。 分层方法一次只预测细一级的占据编码,网络要逐层处理。 在Toys4k的消融中,两级预测变体在体素倒角距离、IoU、网格倒角距离、LPIPS和Uni3D上均优于单级预测变体;在四块A100、批大小为某固定值时,两级变体一个训练步耗时少于单级变体。

推理时重新掩码最早揭开的单元,让它们在细层单元出现后被重新预测。 同一步中揭开的单元是各自独立采样的,可能彼此不一致;已有采样器或不再掩码,或随机、按置信度选择重掩码对象。 在图像条件生成的采样器对比中,该策略的体素倒角距离为0.01107、IoU为0.4158,优于MDLM的0.01162与0.4035、ReMDM的0.01115与0.4085、ReMDM-conf的0.01116与0.4139以及MaskGIT的0.01243与0.3987。

启示与展望

演示范围是单物体体素生成:图像条件任务在TRELLIS-500K上训练、在训练中未见的Toys4k对象上测试,无条件任务使用ShapeNet的椅子、飞机、桌子和步枪四类,形状以最细层级的表面占据表示。作者提出下一步把框架用于包含多个物体的大场景和长序列三维形状。

作者指出训练集规模是限制,约用K个资产训练,不到TRELLIS.2所用约K个资产的四分之一,并预期更多数据会进一步提升。作者同时说明,与其它掩码离散流匹配模型一样,一步内独立揭开多个占据编码可能彼此不一致,现有策略只是减轻这种不一致。