Direct3D-S2: Gigascale 3D Generation Made Easy with Spatial Sparse Attention 讲解
一、引言
论文指出,直接从文本或图像生成高质量 3D 模型具有巨大创意潜力,可用于虚拟世界构建、产品原型设计等众多实际应用。然而,使用体素表示(如符号距离函数 SDF)生成高分辨率 3D 形状面临巨大的计算和内存挑战。针对此问题,研究者们提出了 Direct3D-S2,这是一个基于稀疏体素的可扩展 3D 生成框架,通过Spatial Sparse Attention(SSA)机制显著提升了生成质量和训练效率。
二、相关工作
-
多视图生成与 3D 重建:现有方法多利用多视图扩散模型结合 2D 图像先验模型生成 3D 形状的多视图图像,再进行 3D 重建。但这些方法在保持多视图一致性与形状质量方面存在困难,且依赖于渲染监督,增加了训练复杂度与计算成本。
-
大规模 3D 潜在扩散模型:分为基于 vecset 的隐式方法和基于体素的显式方法。隐式方法虽能生成高质量 3D 形状,但受限于 vecset 大小,训练效率低。显式方法具有较好可解释性与训练 simplicity,但受限于 GPU 内存需求与注意力机制计算成本,在高分辨率输出方面存在局限。
-
高效大规模 token 生成:Native Sparse Attention(NSA)通过自适应 token 压缩降低注意力计算 token 数量,已在大型语言模型和视频生成中取得良好效果。本文将 token 压缩扩展到 3D 数据,提出 SSA 机制,以适应 3D 数据特性。
三、方法
-
稀疏 SDF VAE:论文提出了一种端到端的稀疏 SDF VAE(SS-VAE),采用对称编码器 - 解码器网络架构,利用稀疏 3D 卷积网络和变换器网络的混合框架,将高分辨率稀疏 SDF 体素编码为稀疏潜在表示,再解码重建 SDF 体积。通过聚焦于绝对 SDF 值低于阈值的稀疏体素,有效降低了计算成本。
-
Spatial Sparse Attention(SSA):针对 SS-VAE 分辨率提升导致输入 tokens 长度大幅增加、注意力计算效率低下的问题,论文提出 SSA 机制。该机制基于 3D 坐标对 key 和 value tokens 进行空间相干块划分,并进行块内 token 选择以加速计算。SSA 包括稀疏 3D 压缩、空间块选择和稀疏 3D 窗口三个核心模块,通过这些模块构建对应的 key-value 对,进行注意力计算,并根据门控分数加权聚合结果。
-
稀疏条件机制:为减少背景区域带来的计算开销与对生成网格与条件图像对齐的不利影响,论文提出稀疏条件机制,选择性地从输入图像中提取和处理稀疏前景 tokens,用于交叉注意力计算。
-
校正流:论文采用校正流目标训练生成模型,定义前向过程为数据分布与标准正态分布之间的线性轨迹,生成模型学习预测从噪声样本到数据分布的速度场。
四、实验
-
数据集:Direct3D-S2 在 Objaverse、Objaverse-XL 和 ShapeNet 等公开可用的 3D 数据集上进行训练,经过严格筛选后获得约 452k 个 3D 资产用于训练。对于图像条件的 DiT 训练,每个网格渲染 45 个 RGB 图像。为评估生成网格的几何保真度,建立了包含来自专业社区的高度详细图像的基准测试,并采用 ULIP-2、Uni3D 和 OpenShape 等多模态模型进行定量评估。
-
实现细节:详细介绍了 VAE 和 DiT 的训练配置,包括输入 SDF 值阈值、编码器下采样因子、潜在表示通道维度、不同损失权重、优化器选择与学习率设置等。同时,阐述了 DiT 的层数、隐藏维度、组查询注意力配置、SSA 机制参数设置、条件图像特征提取方式、输入图像分辨率以及逐步提升训练分辨率的策略等。
-
定量与定性比较:通过与 Trellis、Hunyuan3D 2.0、TripoSG、Hi3DGen 等先进图像到 3D 方法的对比,验证了 Direct3D-S2 框架的有效性。结果表明,Direct3D-S2 在三个评估指标上均优于其他方法,生成的网格与输入图像具有更好的一致性。定性比较显示,其他方法因分辨率限制难以捕捉精细结构,而 Direct3D-S2 能生成高分辨率网格,即使在复杂细节上也能取得优越质量。此外,用户研究也表明 Direct3D-S2 在图像一致性和整体几何质量方面具有统计学优势。
-
VAE 比较:SS-VAE 在包含复杂几何结构的验证集上展现出优越的重建精度,且在 512³ 和 1024³ 分辨率下均表现出显著的性能提升。同时,SS-VAE 的训练效率更高,仅需 8 个 A100 GPU 训练 2 天,相较于其他方法所需的至少 32 个 GPU 有明显优势。
-
消融研究:研究了不同分辨率下 Direct3D-S2 的生成结果,发现随着分辨率提升,网格质量逐步改善。在 512³ 分辨率下,验证了 SSA 的三个模块的效果,结果表明各模块协同作用可显著提升网格质量。此外,比较了不同注意力机制的运行时间,SSA 在 tokens 数量多时速度优势明显。还验证了稀疏条件机制的有效性,表明其可使生成网格与输入图像更好地对齐。
五、结论与局限性
论文提出的 Direct3D-S2 框架通过 SSA 机制显著加速了 DiT 的训练和推理速度,端到端对称稀疏 SDF VAE 的集成进一步增强了训练的稳定性和效率。大量实验证明,Direct3D-S2 在生成质量上优于现有先进图像到 3D 方法,且仅需 8 个 GPU 即可进行训练。然而,SSA 的前向传递加速比明显小于后向传递,主要是由于前向传递中 top-k 排序操作引入的计算开销。作者表示将在未来工作中优先优化这些操作。