Flash3D:通过联合硬件-几何局部性实现点转换器的超扩展
计算机视觉与模式识别
2024-12-24 v1
摘要
近期研究认识到规模在 3D 学习中的强大作用(例如 PTv3)以及注意力机制(例如 FlashAttention)。然而,当前的点云骨干网络未能在一个视角中实现几何局部性、注意力机制和 GPU 架构的全局统一。本文引入 Flash3D 转换器,通过基于完美空间哈希(PSH)的原则性局部机制实现几何局部性与 GPU tiling 的对齐。与 GPU tiling 的常见对齐自然地将我们的 PSH 局部机制与 FlashAttention 融合在一起,几乎没有额外开销。该机制为骨干网络中的灵活设计选择提供了便利,结果在下游任务方面表现优异。Flash3D 在基准数据集上超过最先进的 PTv3 结果,实现了 2.25 倍速度提升和 2.4 倍内存效率提升。这种效率使我们能够在不增加额外开销的前提下扩展到更广的注意力范围和更大的模型规模。这种规模化使 Flash3D 能够在相同的计算预算下实现比 PTv3 更高的任务准确率。
引用
@article{arxiv.2412.16481,
title = {Flash3D: Super-scaling Point Transformers through Joint Hardware-Geometry Locality},
author = {Liyan Chen and Gregory P. Meyer and Zaiwei Zhang and Eric M. Wolff and Paul Vernaza},
journal= {arXiv preprint arXiv:2412.16481},
year = {2024}
}