论文
将并行序列模型扩展到基础规模视觉编码器
Scaling Parallel Sequence Models to Foundation-Scale Vision Encoders
摘要
视觉基础模型受到自注意力二次成本的瓶颈,这限制了可用的分辨率并增加了大规模预训练的成本。线性注意力和状态空间模型等次二次替代方案降低了这种成本,但通常将图像序列化为一维词元流并削弱对视觉重要的二维空间结构。广义空间传播网络(GSPN)通过线扫描循环直接在 2D 网格上传播上下文,无需位置嵌入即可实现近线性的复杂性,但很少用作基础规模编码器。我们提出了 C-GSPN,一种基于 2D 空间传播的基础规模视觉编码器。 C-GSPN 通过三项改进使算子变得实用:(1) 快速 GSPN CUDA 内核,将每步启动融合到单个 warp 专用实现中,具有共享内存平铺、合并访问和紧凑的多通道传播,达到峰值内存带宽的 90% 以上,运行速度比原始 GSPN 实现快 40--52 倍; (2)具有融合归一化的压缩潜在空间传播块,将内核级速度转化为块级和模型级效率; (3) 两阶段跨算子 蒸馏 方案,由注意力教师训练新架构,无需从头开始进行基础规模训练。 C-GSPN 采用 6 亿图像-文本对进行提炼,以减少 15% 的参数来匹配同构 ViT 基线,将 ADE20K 分割提高 +2.1%,从头开始使用所需数据的一小部分传输到高分辨率,并通过单通道、无平铺推理在 2K 下提供 4 倍的端到端块加速。