[论文解读] LighTN: Light-weight Transformer Network for Performance-overhead Tradeoff in Point Cloud Downsampling

Xu Wang, Yi Jin+4 more|arXiv (Cornell University)|Feb 13, 2022

Optical Imaging and Spectroscopy Techniques被引用 5

一句话总结

LighTN 提出了一种用于面向任务的点云下采样轻量级 Transformer 网络，通过使用单头自相关模块替代传统的 QKV 投影，减少了计算开销，同时保持了高性能。它通过一种新颖的采样损失函数和扩展-缩减前馈网络缩放策略，在分类和配准任务中实现了最先进性能，且资源消耗极低。

ABSTRACT

Compared with traditional task-irrelevant downsampling methods, task-oriented neural networks have shown improved performance in point cloud downsampling range. Recently, Transformer family of networks has shown a more powerful learning capacity in visual tasks. However, Transformer-based architectures potentially consume too many resources which are usually worthless for low overhead task networks in downsampling range. This paper proposes a novel light-weight Transformer network (LighTN) for task-oriented point cloud downsampling, as an end-to-end and plug-and-play solution. In LighTN, a single-head self-correlation module is presented to extract refined global contextual features, where three projection matrices are simultaneously eliminated to save resource overhead, and the output of symmetric matrix satisfies the permutation invariant. Then, we design a novel downsampling loss function to guide LighTN focuses on critical point cloud regions with more uniform distribution and prominent points coverage. Furthermore, We introduce a feed-forward network scaling mechanism to enhance the learnable capacity of LighTN according to the expand-reduce strategy. The result of extensive experiments on classification and registration tasks demonstrates LighTN can achieve state-of-the-art performance with limited resource overhead.

研究动机与目标

• 为解决在低功耗应用中基于 Transformer 的点云下采样资源开销过高的问题。
• 开发一种即插即用、端到端的下采样解决方案，兼容预训练的任务网络。
• 通过消除 QKV 投影矩阵并利用扩展-缩减缩放优化网络深度，实现性能与效率的平衡。
• 通过一种新颖的采样损失函数，引导网络聚焦于关键且分布均匀的点。

提出的方法

• 提出一种单头自相关模块，无需单独的查询、键和值投影矩阵即可计算注意力，从而降低参数量和 FLOPs 开销。
• 采用对称矩阵计算，天然满足排列不变性，提升了稳定性和效率。
• 设计一种多组件采样损失，结合 Chamfer 距离、替换损失和带指数温度缩放的软分配损失，以提升点的覆盖度和均匀性。
• 在前馈网络中应用扩展-缩减策略，以极低的 FLOPs 和参数增加实现深度扩展。
• 在自相关模块中完全移除所有投影矩阵，实现了性能与资源成本之间的最优权衡。
• 默认配置采用第二层扩展比为 2 的三层 FFN，兼顾准确率与效率。

实验结果

研究问题

• RQ1 轻量级 Transformer 架构是否能在显著降低计算与内存开销的前提下，实现点云下采样任务的最先进性能？
• RQ2 移除 QKV 投影矩阵对下采样任务中自注意力机制的性能与效率有何影响？
• RQ3 在轻量级下采样框架中，前馈网络的最优配置是什么，以实现模型深度与资源成本的平衡？
• RQ4 所提出的采样损失函数在引导网络保留关键且分布均匀的点方面效果如何？
• RQ5 自相关机制在下采样任务中是否在判别能力与效率方面优于标准点积注意力？

主要发现

• LighTN 在 ModelNet40 分类和 ShapeNet 配准基准上实现了最先进性能，且资源开销显著降低。
• 与三头变体相比，完全移除所有 QKV 投影矩阵可将 FLOPs 降低高达 65.6%，参数量减少 3.6%，且准确率下降可忽略不计。
• 采用对称矩阵计算的自相关模块在判别能力方面表现更优，且开销更低，优于标准单头点积注意力。
• 在软分配损失组件中使用指数温度函数相比二次缩放，能提升模型收敛速度与准确率。
• 采用 L=3 和 r=2 的扩展-缩减策略，平均仅增加 0.0087G FLOPs 和 0.0488M 参数，但显著提升了性能。
• 结合 L_soft(exp(t)) 和 L_repl 组件的所提采样损失函数实现了最佳准确率，证明了其在引导点选择方面的有效性。