MoCA: Mixture-of-Components Attention for Scalable Compositional 3D Generation¶
会议: NeurIPS 2026
arXiv: 2512.07628
项目主页: https://lizhiqi49.github.io/MoCA
领域: 3D生成与编辑
关键词: 组合式3D生成, 稀疏注意力机制, 扩散Transformer, 混合组件注意力, 部件解耦
一句话总结¶
针对组合式3D生成中全局注意力随部件数量暴涨导致的二次方计算瓶颈,本文提出MoCA,通过轻量级重要性路由选择Top-k关键组件保留全量tokens、并对远端不重要组件进行压缩降维,首次将端到端组合式生成扩展至单样本32个细粒度部件与复杂场景实例。
研究背景与动机¶
组合性(Compositionality)是三维资产设计与构建流程的核心基石。无论是单个复杂物体由多个语义明确的部件组装而成,还是大尺度室内外场景由多个独立物体摆放构成,解耦的部件级与实例级表示都是支持部件复用、局部精细化编辑、动力学仿真和物理交互所必需的前提。然而,近期的三维扩散Transformer(DiT)模型多将整个物体或场景视为不可分割的单一实体进行表征与去噪,生成的隐空间缺乏结构化语义解耦,使得下游部件编辑或材质定制极为受限。
为解决这一问题,前沿探索尝试将预训练3D生成先验扩展到多组件联合生成框架中。但这类方案在面对细粒度或大尺度资产扩展时面临严重的计算瓶颈:现有方法在跨组件全局交互中,普遍将所有组件的隐空间tokens直接拼接并计算全局注意力。若每个资产包含 \(N\) 个组件、每个组件由 \(L\) 个隐向量表示,注意力序列长度即达 \(N \times L\),全局注意力计算复杂度呈二次方 \(O(N^2 L^2)\) 爆炸。此外,空间上并非所有组件间都存在强相互作用——例如建模角色的手指主要依赖手腕与小臂,而椅子在场景中的位姿主要与临近的桌子相关。盲目让所有tokens两两进行密集交互,不仅浪费了宝贵的算力与显存,更将可建模组件数量严重限制在极小规模(通常部件 \(\le 8\)、场景实例 \(< 20\))。
本文的切入角度是引入类似于混合专家模型(MoE)的稀疏路由与信息分层压缩思路,重构跨组件全局注意力层。核心 idea:通过基于锚点重要性的轻量级组件路由,动态挑选Top-k关键交互组件保留全量tokens密集通信,同时将远端其余非关键组件压缩为紧凑隐向量作为全局几何与布局先验,在大幅降低注意力上下文长度的同时实现兼顾全局上下文与局部几何一致性的可扩展组合式3D生成。
方法详解¶
整体框架¶
MoCA 是一个面向组合式 3D 物体与场景的扩散 Transformer(DiT)生成系统。系统以单张物体图像或带有实例掩码的场景图像为条件,通过整流流动匹配(Rectified Flow Matching)联合去噪生成 \(N\) 个组件的清洁 vecset 隐向量 \(\mathcal{Z} = \{\mathbf{z}_i\}_{i=1}^N \in \mathbb{R}^{N \times L \times D}\)。去噪完成后,每个组件的隐变量输入冻结的预训练 vecset 解码器中解码出各自的连续空间场,再通过 Marching Cubes 算法提取独立的三角形网格,合并构成整体资产。
整个 DiT 骨干网络由局部注意力模块(Local Attention Block)与混合组件注意力模块(Mixture-of-Components Attention, MoC Block)交替堆叠而成。为了在全局块中实现稀疏高效的跨组件通信,MoCA 在输入端为每个组件引入可学习的压缩查询 tokens 与代表全局语义的单一锚点 token,使得全局上下文长度从朴素拼接的 \(N \times L\) 压缩为 \(L + kL + (N - k - 1)\frac{L}{\sigma}\)(其中 \(\sigma\) 为压缩比率,默认取 8)。
%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
A["单图/多视条件 + 随机噪声隐变量"] --> B["阶段1:可学习Token打包与组件特征压缩<br/>跨注意力提取压缩tokens与单锚点token"]
B --> C["阶段2:局部注意力特征更新<br/>因果掩码阻断vecset泄露并精炼局部表示"]
C --> D["阶段3:基于重要性的组件路由<br/>锚点点积打分与无辅助损失随机负载均衡"]
D --> E["阶段4:混合组件注意力融合<br/>Top-k全量+远端压缩tokens加权聚合"]
E --> F["冻结Vecset解码器独立重建<br/>输出解耦3D部件与完整网格"]
关键设计¶
1. 可学习Token打包与组件特征压缩:为无序几何隐空间构建紧凑上下文载体
针对三维点云与无序隐向量集合(vecset)无法直接应用传统二维卷积或池化降维的特质,本文为每个组件的加噪隐向量额外引入 \(N_p = \frac{L}{\sigma}\) 个可学习压缩查询 tokens \(\mathbf{p} \in \mathbb{R}^{N_p \times D}\),以及一个代表该组件全局紧凑特征的单锚点 token \(\bar{\mathbf{p}} \in \mathbb{R}^{1 \times D}\)。在送入 DiT 之前,通过跨注意力层(Cross-Attention)对加噪 vecset 隐变量 \(\mathbf{z}_i\) 进行特征打包: $\(\mathbf{p}_i = \text{CrossAttn}(\mathbf{p}, \mathbf{z}_i, \mathbf{z}_i), \quad \bar{\mathbf{p}}_i = \text{CrossAttn}(\bar{\mathbf{p}}, \mathbf{z}_i, \mathbf{z}_i)\)$ 这使得每个组件的输入表示拼接为 \(\mathbf{x}_i = \text{Concat}(\mathbf{z}_i; \mathbf{p}_i; \bar{\mathbf{p}}_i)\),并附加随机 ID 嵌入以消除排列对称性。在随后的局部注意力层中,为了防止新引入的 query tokens 破坏预训练 3D 几何建模先验,vecset tokens \(\mathbf{z}_i\) 被强制限制为仅进行自注意力交互,对附加 tokens 保持盲视;而压缩 tokens \(\mathbf{p}_i\) 则可同时感知 \(\mathbf{z}_i\) 与自身以持续吸收几何更新,锚点 token \(\bar{\mathbf{p}}_i\) 作为最高级抽象可感知该组件的所有 tokens。
2. 基于重要性的组件路由:以注意力机制筛选Top-k关键通信子集
针对空间中并非所有组件都有直接强关联的先验,MoCA 在全局注意力模块中引入轻量级路由机制,计算任意两个组件之间的相对重要性。具体而言,对于第 \(i\) 个组件 \(\mathbf{c}_i\) 与第 \(j\) 个组件 \(\mathbf{c}_j\)(\(j \neq i\)),路由模块将其各自的锚点 token 经由独立线性投影变换为查询向量 \(\bar{Q}_i = \text{Query}(\bar{\mathbf{p}}_i)\) 与键向量 \(\bar{K}_j = \text{Key}(\bar{\mathbf{p}}_j)\),通过点积注意力计算相对重要性权重: $\(o_{i,j} = \text{Sigmoid}\left(\frac{\bar{Q}_i \bar{K}_j^T}{\sqrt{d_{\bar{K}}}}\right)\)$ 这里采用 Sigmoid 替代传统 Softmax 激活函数,有效避免了当组件数量 \(N\) 增大时出现的概率平铺与梯度弥散问题。计算得出全部重要性得分 \(\{o_{i,j}\}_{j \neq i}\) 后,系统选出得分最高的前 \(k\) 个组件保留全量 vecset tokens \(\mathbf{z}_j\),而对其余较远的无关组件仅保留压缩 tokens \(\mathbf{p}_j\),从而兼顾关键连接处的精细几何协同与大局空间排布先验。
3. 混合组件注意力与键向量门控:端到端可导的重加权通信机制
在确定路由分配后,目标组件 \(\mathbf{c}_i\) 需要与其它组件的路由候选集合进行交互。为让组件路由层在反向传播中获得端到端梯度更新,MoCA 将预测的重要性得分 \(o_{i,j}\) 直接作为门控标量,与对应被路由组件的键向量相乘,其查询、键、值向量形式定义如下: $\(\mathcal{Q}_i = \text{Query}(\mathbf{x}_i)\)$ $\(\mathcal{K}_i = \text{Concat}\Big(\text{Key}(\mathbf{z}_i); \, o_{i,1} \cdot \text{Key}(\mathbf{r}_{i,1}); \, \dots; \, o_{i,N} \cdot \text{Key}(\mathbf{r}_{i,N})\Big)\)$ $\(\mathcal{V}_i = \text{Concat}\Big(\text{Value}(\mathbf{z}_i); \, \text{Value}(\mathbf{r}_{i,1}); \, \dots; \, \text{Value}(\mathbf{r}_{i,N})\Big)\)$ 通过在 Softmax 计算前直接调制键向量的振幅,重要性得分能够无损影响注意力权重分布,同时规避了在值向量后乘加权破坏归一化约束导致的数值崩溃。此外,该路由在多头注意力各 head 维度独立执行,赋予不同注意力头学习差异化交互依赖的表达容量。
一个完整示例¶
以生成一台包含 \(N=16\) 个细分部件的复杂工业机器人为例,每个部件分配 \(L=1024\) 个隐空间 tokens,压缩比设定为 \(\sigma=8\)(即压缩 tokens 为 128 个),重要组件选取比例为 25%(即 \(k=4\))。 1. 输入打包:机械臂手爪部件(记为 \(\mathbf{c}_1\))拥有 1024 个 vecset tokens、128 个压缩 tokens 和 1 个单锚点 token。 2. 局部更新:在局部块中,1024 个 vecset tokens 仅在内部自注意力更新几何轮廓;128 个压缩 tokens 跨步关注 vecset 并形成压缩概要;单锚点聚合整爪全局状态。 3. 重要性路由:锚点 token 与其余 15 个部件的锚点点积打分。路由模块准确识别出紧邻的手腕、前臂、指尖关节和传感器(\(k=4\) 个部件)得分最高。 4. MoC 注意力计算:手爪部件在全局交互时,对这 4 个相邻关键部件接入全部 \(4 \times 1024 = 4096\) 个原始 tokens,进行精细接缝与物理连接对齐;对基座、履带、外壳等其余 11 个远端部件,仅分别接入 128 个压缩 tokens(共 \(11 \times 128 = 1408\) 个 tokens),提供整体比例与全局布局约束。 5. 计算收益:全局交互的上下文长度仅为 \(1024 + 4096 + 1408 = 6528\),较朴素全注意力拼接的 \(16 \times 1024 = 16384\) 缩减了超过 60%,在保证活动关节无缝拼合的同时消除了显存爆炸。
损失函数 / 训练策略¶
MoCA 采用整流流动匹配(Rectified Flow Matching)目标进行监督。设所有组件的干净 vecset 隐向量集合为 \(\mathcal{Z}_0 = \{\mathbf{z}_i^0\}_{i=1}^N \in \mathbb{R}^{N \times L \times D}\),在时间步 \(t \sim \mathcal{U}(0, 1)\) 下沿线性轨迹混合高斯噪声 \(\epsilon \sim \mathcal{N}(0, \mathbf{I})\): $\(\mathcal{Z}_t = (1 - t)\mathcal{Z}_0 + t\epsilon\)$ 网络训练用于预测速度场 \(v_\theta(\mathcal{Z}_t, t, \mathcal{I})\),优化均方误差目标: $\(\mathcal{L} = \mathbb{E}_{\mathcal{Z}, \epsilon, t} \left[ \| (\epsilon - \mathcal{Z}_0) - \mathbf{v}_\theta(\mathcal{Z}_t, t, \mathcal{I}) \|^2 \right]\)$ 无辅助损失的随机负载均衡(Auxiliary-Loss-Free Load Balancing):为防止路由器在训练初期陷入仅反复选择少数特定组件、导致其他组件欠拟合或陷入局部极小的问题,MoCA 舍弃了引入复杂超参且梯度不稳定的辅助损失,转而在训练阶段将归一化的重要性 logits 转换为概率分布,从中随机采样 Top-k 组件索引以促进充分探索;而在推理阶段严格恢复为确定性 Top-k 选取,兼顾训练多样性与生成稳定性。
在具体训练中,物体模型首先在 32 张 H20 GPU 上以 \(L=512\)、batch size=1024 预训练 40K 步,随后在 \(L=1024\) 下微调 10K 步;场景模型则以 4 张 isometric 渲染图与实例掩码作为 7 通道条件,在 3D-FRONT 数据集上用 LoRA 微调 10K 步。
实验关键数据¶
主实验¶
在部件级生成任务中,模型在包含 200 个物体的 PartObjaverse-Tiny 与随机抽取的 100 个 ABO 资产上进行几何评测。评测采用熔融表面点云计算 Chamfer Distance(CD,越低越好)、双阈值 F-score(Fscore-0.1 与 Fscore-0.05,越高越好),以及衡量部件重叠侵入程度的 Self-IoU(越低越好)。
| 数据集 | 指标 | MoCA (本文) | PartPacker | PartCrafter* | HoloPart |
|---|---|---|---|---|---|
| PartObjaverse | Self-IoU ↓ | 0.0125 | 0.0120 | 0.0224 | 0.0142 |
| PartObjaverse | CD ↓ | 0.1010 | 0.1105 | 0.1195 | 0.1145 |
| PartObjaverse | Fscore-0.1 ↑ | 0.8708 | 0.8484 | 0.8169 | 0.8340 |
| PartObjaverse | Fscore-0.05 ↑ | 0.6882 | 0.6510 | 0.6236 | 0.6671 |
| ABO | Self-IoU ↓ | 0.0116 | 0.0119 | 0.0136 | 0.0139 |
| ABO | CD ↓ | 0.1027 | 0.1094 | 0.1124 | 0.1168 |
| ABO | Fscore-0.1 ↑ | 0.8755 | 0.8646 | 0.8515 | 0.8523 |
| ABO | Fscore-0.05 ↑ | 0.6871 | 0.6801 | 0.6759 | 0.6772 |
在场景级生成任务中,在 3D-FRONT 测试集上评估场景与实例级别的 CD 及 F-score:
| 方法 | Self-IoU ↓ | CD-Scene ↓ | Fscore-Scene ↑ | CD-Obj ↓ | Fscore-Obj ↑ |
|---|---|---|---|---|---|
| MIDI | 0.0011 | 0.1548 | 0.8238 | 0.1442 | 0.7996 |
| PartCrafter | 0.0036 | 0.1366 | 0.8379 | - | - |
| MoCA (本文) | 0.0006 | 0.1201 | 0.8637 | 0.1163 | 0.8380 |
消融实验¶
在 Trellis-500K 数据集上以最多 16 个部件训练,并在 PartObjaverse-Tiny 上评测各模块消融结果:
| 配置代号 | 路由全量重要组件 | 压缩保留远端组件 | 门控乘法位置 | 激活函数 | 负载均衡策略 | 多头独立路由 | CD ↓ | Fscore-0.1 ↑ | Fscore-0.05 ↑ |
|---|---|---|---|---|---|---|---|---|---|
| A (无重要路由) | ✗ | ✓ | Key | Sigmoid | ✓ | ✓ | 0.1969 | 0.6604 | 0.4750 |
| B (直接丢弃远端) | ✓ | ✗ | Key | Sigmoid | ✓ | ✓ | 0.1523 | 0.7494 | 0.5465 |
| C (门控乘至Value) | ✓ | ✓ | Value | Sigmoid | ✓ | ✓ | 0.1519 | 0.7612 | 0.5531 |
| D (Softmax激活) | ✓ | ✓ | Key | Softmax | ✓ | ✓ | 0.1451 | 0.7638 | 0.5539 |
| E (无负载均衡) | ✓ | ✓ | Key | Sigmoid | ✗ | ✓ | 0.1361 | 0.7947 | 0.5975 |
| F (单头共享路由) | ✓ | ✓ | Key | Sigmoid | ✓ | ✗ | 0.1212 | 0.8215 | 0.6178 |
| G (完整模型) | ✓ | ✓ | Key | Sigmoid | ✓ | ✓ | 0.1180 | 0.8259 | 0.6233 |
在推理时间分解方面,当部件扩展至 \(N=32\) 时,MoCA 的全局注意力计算耗时仅为 79 ms,相较 PartCrafter 的 142 ms 降低了近 45% 的计算负担。
关键发现¶
- 关键组件全量交互与远端上下文压缩缺一不可:消融实验 A 显示,完全放弃 Top-k 全量路由而仅依赖压缩特征,CD 指标急剧恶化至 0.1969(大幅暴跌 66.8%),证明精细几何拼合强依赖局部高分辨率交互;消融实验 B 证实若直接抛弃未选中的远端组件,CD 同样劣化至 0.1523,证明粗粒度全局空间先验是避免组件脱节或定位漂移的关键。
- 门控加权位置与激活函数选择至关重要:配置 C 将路由权重乘在 Value 向量上导致 CD 显著升高至 0.1519,原因在于后 Softmax 调制打破了注意力分布总和为 1 的归一化约束;配置 D 证明在组件数较大时 Softmax 导致门控值过小产生梯度弥散,而 Sigmoid 能稳定激活。
- 随机路由有效破解 MoE 崩溃陷阱:在训练中引入无辅助损失的随机概率采样(配置 E vs G)使 Fscore-0.1 从 0.7947 提升至 0.8259,验证了早期探索阶段平衡各组件注意力负载的有效性。
亮点与洞察¶
- 信息分层而非简单二值丢弃:将注意力对象划分为“局部全分辨率”与“远端压缩概要”,既打破了注意力平方复杂度魔咒,又保留了全局几何与语义锚定,避免了纯局部窗口注意力容易丢失全局上下文的缺陷。
- 无辅助损失的免调参负载均衡机制:借鉴大语言模型 MoE 训练的前沿经验,通过在训练期将 logits 转换为采样概率进行随机路由、推理期确定性贪心 Top-k,以零额外超参数代价解决了组件路由器崩溃问题。
- 统一物体部件与场景实例的架构范式:MoCA 将 Part-level 3D 物体生成与 Instance-level 3D 场景生成统一在同一套混合组件注意力范式下,最大组件数量支持翻倍至 32 个,为通用组合式 3D 基础模型提供了可复用的架构蓝图。
局限与展望¶
- 冻结 VAE 对微小部件的重建瓶颈:当前方法沿用冻结的全局 vecset VAE,当复杂资产被细分为多达 32 个部件时,极小体积部件在全局坐标空间中所占体素或采样点极少,可能带来隐空间重建失真;未来可探索针对部件级局部坐标系自适应微调的 VAE 结构。
- PyTorch 纯原生路由索引带来的工程开销:如运行时间分析所示,虽然全局注意力计算时间减半,但由于基于 Python/PyTorch 的排序与切片操作耗时在 \(N=32\) 时达到了 36 ms,未来需要编写定制化的 Triton/CUDA 算子融合核函数以释放极致吞吐潜力。
- 多视角与物理先验结合:目前的场景生成高度依赖实例掩码先验,未来可探索将大语言模型的物理常识布局先验内化至路由层,实现纯文本直接端到端驱动的大规模复杂场景合成。
相关工作与启发¶
- vs PartCrafter (Lin et al., 2025):PartCrafter 采用朴素的全部部件 tokens 直接拼接计算全密集注意力,计算量呈 \(O(N^2 L^2)\) 增长,因而受限于极少量粗粒度部件(通常 \(<8\));MoCA 引入重要性路由与远端压缩,将计算复杂度大幅收敛,首次在单样本内稳定支持 32 个组件的精细解耦。
- vs MIDI (Huang et al., 2025):MIDI 同样通过密集注意力生成多实例场景,但难以扩展至 20 个以上的复杂场景且易出现部件混淆与表面断裂;MoCA 结合 Sigmoid 锚点门控与局部-全局混合更新,在场景与实例级 CD 上均取得了显著领先。
- vs HoloPart / PartPacker (Yang et al., 2025; Tang et al., 2025):HoloPart 采用“先整体生成、后分割补全”的多阶段流程,易产生部件间接缝与穿模;PartPacker 无法控制生成部件数量且分割粒度粗糙。MoCA 在 end-to-end 框架下保证了严格的部件间解耦与极低的 Self-IoU。
评分¶
- 新颖性: ⭐⭐⭐⭐⭐ 巧妙将 MoE 稀疏路由与 token 压缩机制迁移至 3D 组合式扩散 Transformer,架构设计优雅且切中痛点。
- 实验充分度: ⭐⭐⭐⭐⭐ 涵盖物体部件与大尺度场景双任务,包含详细消融、真实图像泛化测试与细致的算力/时间分解。
- 写作质量: ⭐⭐⭐⭐⭐ 问题表述清晰,动机阐述层层递进,公式推导严谨且实验结论说服力极强。
- 价值: ⭐⭐⭐⭐⭐ 突破了组合式 3D 生成在多组件规模上的扩展瓶颈,对数字孪生、游戏资产自动化生成具有重大实用价值。