LRM-典型 Transformer 在视觉领域的应用,单个图像生成3D图像
https://yiconghong.me/LRM
一、Abstract
第一个大型重建模型(LRM),它可以在5秒内从单个输入图像预测物体的3D模型。LRM采用了高度可扩展的基于transformer的架构,具有5亿个可学习参数,可以直接从输入图像中预测神经辐射场(NeRF)。
算法具体步骤
- 图像编码(Image Encoder)
- 以 RGB 图像为输入,采用预训练的视觉 Transformer(ViT)模型 DINO 将图像编码为 patch 级特征令牌,其中为 patch 总数,为编码器潜在维度。此步骤利用 DINO 学习图像显著内容结构和纹理的可解释注意力,相比其他语义导向表示,能更好保留重建 3D 空间几何与颜色所需的详细结构和纹理信息,且不仅使用聚合特征的 [CLS] 令牌,还利用整个特征序列。
- 图像到三平面解码(Image-to-Triplane Decoder)
- 相机特征构建与嵌入:将输入图像的 4x4 相机外参矩阵展平并与相机焦距和主点拼接构建相机特征,经相似变换归一化使输入相机对齐同一轴,再通过多层感知器(MLP)映射为高维相机嵌入,同时归一化相机内参(焦距和主点)。
- 三平面表示获取:以可学习的空间位置嵌入为导向,通过跨注意力将图像和相机特征投影到三平面表示,其包含三个轴对齐特征平面、和,每个平面维度。在正向传播中,基于相机特征和图像特征,图像到三平面变压器解码器的各层通过调制和跨注意力分别更新初始位置嵌入为最终三平面特征,包括相机特征调制(基于 DiT 自适应层归一化思想实现)、变压器层内的跨注意力、自我注意力和多层感知器子层操作,最终经上采样和重塑得到三平面表示。
- 三平面 - NeRF(Triplane-NeRF):利用从三平面表示查询的点特征,通过包含多个线性层(ReLU 激活)的预测 RGB 和密度,输出维度为 4(前三维为 RGB 颜色,最后一维为场密度)用于体绘制。
- 训练目标(Training Objectives):从单输入图像生成 3D 形状,训练时利用额外侧视图指导重建。对训练数据中每个形状,考虑个随机选择的侧视图进行监督,通过最小化渲染视图与真实视图间的图像重建目标(包括归一化像素级损失和感知图像块相似度加权和)优化模型。
LRM 包含一个图像编码器,它将输入图像编码为逐块特征标记
二、方法
LRM 的整体架构,是一个基于transformer、完全可微的编解码器框架,用于单图像到神经辐射场(NeRF)的重建。LRM 应用一个预训练的视觉模型(DINO)对输入图像进行编码(第 1 节),在此过程中,图像特征通过一个大型变换器解码器借助交叉注意力被投影到三维三平面表示中(第 2 节),随后利用一个多层感知器来预测用于体积渲染的点颜色和密度(第 3 节)。整个网络基于大约一百万个三维数据进行端到端训练,并采用简单的图像重建损失函数。
1. IMAGE ENCODER(图像编码器)
使用一个经过自蒸馏训练的模型DINO,用它来学习对图像中显著内容的结构和纹理。相比其他面向语义的表示等视觉特征ImageNet预训练ResNet或CLIP,精细的结构和纹理信息更重要,因为后续LRM要使用它来重建三维空间的几何和颜色。
PS:
预训练的视觉模型 DINO 是一种基于自监督学习的先进模型,以下是简单解释:
- 核心原理:DINO 采用对比学习和自蒸馏技术,通过对输入图像进行多尺度裁剪、变换等操作生成不同视图,让模型学习将同一图像的不同视图拉近,不同图像的视图推远,以此学习判别性特征,同时利用自蒸馏将教师模型的知识传递给学生模型来提升性能.
- 架构基础:通常建立在视觉 Transformer(ViT)架构之上,将图像分割成图像块并转化为线性嵌入序列,借助自注意力机制捕捉图像块间的长期依赖关系,从而有效获取图像的全局和局部信息.
2. IMAGE-TO-TRIPLANE DECODER(图像到三平面解码器)
实现了一个基于transformer的解码器,将图像和编码器捕捉的特征投射到可学习的空间位置编码上,并将它们转换为三平面表示。这个解码器是一个先验网络,用大规模的数据进行训练,以提供必要的几何和外观信息,来补偿单图像重建的模糊性。
2.1 相机特征
把代表相机和外面世界位置关系的一个矩阵(4×4 的相机外参矩阵)给弄平了,再和相机的焦距、画面中心这些参数组合在一起,就形成了相机特征。然后呢,还要对这些参数做个处理(归一化),让所有相机好像都在同一个 “起跑线” 上(在同一轴上对齐),这样方便后面进行计算,而且真实的这些参数在训练模型的时候才会用到哦,等模型训练好了实际用的时候就不用啦。之后呢,再用一个多层感知器把整理好的相机特征变得维度更高一些,这样就得到了更方便后面用的相机嵌入
2.2 三平面表示
对于三平面表示,以特定方式构建并利用空间位置嵌入等逐步更新完善,其中交叉注意力模块建立图像与三平面联系,自注意力模块梳理三平面内关系,多层感知器子层再处理,如同层层加工的流水线,共同塑造准确的三平面特征,为 3D 重建夯实基础。
相机特征调制
按照一种特定的计算规则(参考那个公式),根据相机特征来改变图像特征的一些 “数值”,让它更符合我们想要的样子,而且这个调整在每个处理图像和特征的小环节(注意力子层)里都会做。
变换器层
这里面每个变换器层就像一个小 “加工站”,里面有三个小 “工序”(交叉注意力子层、自注意力子层、多层感知器子层),每个工序处理的东西(输入标记)都会先按照相机特征来做个 “预处理”。第一个工序(交叉注意力子层)呢,就是把代表三维物体样子的特征(三平面隐藏特征)和图像特征联系起来,让它们能 “对上号”,不过我们没提前规定它们具体怎么对应的,就让模型自己去学怎么把二维的图像和三维的特征对应好。然后第二个工序(自注意力子层)呢,就是在已经联系好的这些三维特征里,再找找它们自己内部有啥关系,让它们之间更有条理。最后第三个工序(多层感知器子层)再按照常规的做法处理一下。经过这一套,每个小 “加工站” 处理完的结果就传给下一个 “加工站” 接着处理。 最终输出由一个可学习的反卷积层上采样并重塑为最终的三平面表示T。
3. TRIPLANE-NERF(三平面 - 神经辐射场)
三、实验结果
下图可视化了一些从单个图像中进行重建的示例(包括在各种数据集下,具有不同纹理的不同主题的真实、生成和渲染图像)。示例中不仅复杂的几何模型被正确地建模(如椅子、盘子和擦拭),而且模型最终也保存了物体图片上的高频细节,如木孔雀的纹理。从长颈鹿、企鹅和熊等例子。并且我们能够看到,LRM可推断出语义上合理的形状遮挡部分,这意味着已经学习了有效的交叉形状先验,都反映了该模型的强大的泛化能力。
研究人员将LRM与通过使用二维扩散模型生成多视图图像的One-2-3-45进行了比较,可以看到,LRM产生了更清晰的细节和一致的表面。在图的最后一行中,两个例子测试了One-2-3-45显示了更糟糕的重建结果。
四、创新
- 大规模模型与数据训练:构建超 5 亿参数模型,于近百万跨类 3D 数据训练,挖掘通用 3D 先验,提升泛化力,适多种输入图像重建 3D 模型。
- 高效架构设计:基于 Transformer 编解码器与三平面 NeRF,编解码器用跨、自注意力融合 2D 图像和 3D 三平面信息;三平面 NeRF 计算优、利用图像局部性佳,训练靠最小化图像差异,免复杂正则化与调参,适配多数据集。
- 跨模态投影创新:借跨注意力实现 2D 特征至 3D 三平面投影,学习 3D 表示,重建复杂形状与纹理,推断遮挡形状,突破单图像 3D 重建困境。
五、贡献
- 技术提升:于单图像 3D 重建突破,多指标胜前沿法,重建质量优,几何纹理准,助 3D 重建技术用于多行业。
- 实践可行:单 GPU 5 秒重建 3D 形状,含多环节处理,可嵌入多应用,提升 3D 重建应用价值与效率。
- 启发研究:大规模学习理念及架构、训练、多模态创新为后续研究引路,如跨模态学习、模型扩展优化等,促领域发展与智能图像理解。