本研究的作者包括彭高、Jiasen Lu、李宏生、Roozbeh Mottaghi以及Aniruddha Kembhavi。他们分别来自上海人工智能实验室、香港中文大学(CUHK)-商汤联合实验室、华盛顿大学以及艾伦人工智能研究所(AI2)。这篇题为“Container: Context Aggregation Network”的研究论文发表在第三十五届神经信息处理系统会议(NeurIPS 2021)上。
这项研究属于计算机视觉领域,具体聚焦于视觉表示学习的基础网络架构设计。研究的背景源于卷积神经网络(CNN)长期以来在计算机视觉任务中的主导地位,以及近期来自自然语言处理领域的Transformer架构在视觉任务上展现出的强大竞争力。更令人意外的是,研究发现即使是不包含传统卷积或自注意力机制的简单多层感知机(MLP)架构,也能产生有效的视觉表示。尽管CNN、Transformer和MLP-Mixer在表面上被视为截然不同的架构,但本研究旨在揭示它们之间的内在联系。其核心目标是:提供一个统一的视角,表明这些架构实际上是一种更通用的神经网络空间上下文聚合方法的不同特例。基于这一视角,研究者们旨在设计一种新颖的、通用的基础构建模块,它能够兼具Transformer的长程交互能力和CNN的局部归纳偏置所带来的快速收敛优势,从而在图像分类、目标检测、实例分割等多种视觉任务上实现优越的性能。
详细的工作流程分为理论统一、模块设计、网络架构构建以及全面的实验验证四个主要环节。
第一,理论统一与框架构建。研究首先提出了一个广义的上下文聚合模块公式。该公式将神经网络中常见的信息聚合操作抽象为一个核心组件:亲和矩阵(Affinity Matrix)。具体而言,对于一个输入的特征序列,聚合过程可以表示为 y = (A * V) * W1 + x,其中V是输入特征的线性投影,W1是可学习参数,而矩阵A则定义了特征之间信息传播的权重关系,即“亲和力”。A可以是静态的(与输入无关,如卷积核权重),也可以是动态的(由输入计算得出,如自注意力权重)。通过将通道划分为多个“头”(head),可以引入多个亲和矩阵进行多头聚合,以捕获特征空间中不同的关系,提升模型的表示能力。在这个统一框架下,研究者证明了三种主流架构均可被视为特例:Transformer的自注意力机制对应动态生成的亲和矩阵;深度可分离卷积(Depthwise Convolution)对应静态、稀疏且具有局部性和平移不变性的亲和矩阵;MLP-Mixer中的转置MLP操作则对应一个静态、稠密且不共享参数的亲和矩阵。这一理论洞察是后续创新的基石。
第二,新型基础模块的设计。基于上述统一视角,研究者提出了名为“容器”(Container, Context Aggregation Network)的新型基础构建块。Container模块的核心创新在于,它通过可学习的混合系数,将静态亲和矩阵与动态亲和矩阵相结合。其单头形式为:y = ((α * A_dynamic(x) + β * A_static) * V) * W2 + x。其中,A_dynamic(x) 是像Transformer那样由输入动态计算出的亲和矩阵,A_static 是像卷积那样的静态亲和矩阵,α和β是可学习的标量参数,用于自动平衡两种聚合模式的作用。通过设置不同的α、β以及选择特定的静态矩阵(如卷积式或MLP式),Container可以灵活地退化为纯粹的Transformer、深度卷积或MLP-Mixer模块,也能创造出融合二者优势的新变体。在默认配置中,研究者选择卷积式的静态矩阵,以注入对视觉任务有益的局部性和平移不变性偏置。
第三,网络架构的构建。为了进行公平且全面的实验对比,研究者设计了一个分阶段的层级化基础网络架构。该架构包含四个阶段,每个阶段逐步降低空间分辨率(通过块嵌入实现),同时增加通道维度。每个阶段由一系列相同的Container块堆叠而成,每个块内部包含两个子模块:负责空间信息聚合的Container模块(即上述新设计的模块)和负责通道信息融合的前馈网络模块。基于此基础架构,研究者构建了两个主要模型:1)Container:在所有阶段均使用完整的(静态+动态)Container块。2)Container-Light:一种高效变体,在前三个阶段(处理较高分辨率特征图时)关闭计算昂贵的动态亲和矩阵,仅使用静态卷积式聚合;仅在最后阶段(特征图分辨率较低时)启用完整的Container块。这种设计旨在降低高分辨率输入下的计算开销,使其能够高效地应用于目标检测等下游任务。
第四,全面的实验验证与分析。研究进行了广泛而深入的实验,涵盖多个任务和数据集。 图像分类(ImageNet):在ImageNet数据集上,Container(22.1M参数)取得了82.7%的Top-1准确率,比参数量相当的DeiT-S高出2.8个百分点。Container-Light(20.0M参数)也达到了82.0%的准确率,且计算量更低、吞吐量更高。实验还包括收敛速度分析,结果显示Container结合了CNN的快速收敛和Transformer的强大表示能力,仅用200轮训练就达到了DeiT-S在300轮才能达到的79.9%准确率。数据效率实验表明,在仅使用10%ImageNet数据训练时,Container-Light的准确率(61.8%)远超DeiT-S(39.3%),显示出其强大的归纳偏置和样本效率。此外,研究者还在Container框架内轻松复现并扩展了多种已有架构(如分层DeiT、多头MLP-Mixer),验证了框架的通用性和灵活性。可视化分析发现,在网络早期层学习到的静态亲和矩阵中,自发地出现了类似卷积的局部模式,这为模型的行为提供了有趣的解释。 目标检测与实例分割(COCO):使用RetinaNet和Mask R-CNN检测器,并以Container-Light作为主干网络进行实验。在参数量和计算量可比的情况下,Container-Light显著优于传统的ResNet-50主干网络。例如,在RetinaNet上,mAP从36.5%提升至43.8%(+7.3);在Mask R-CNN上,检测框mAP从38.2%提升至45.1%(+6.9),实例掩码mAP从34.7%提升至41.3%(+6.6)。Container-Light也超越了同期其他基于Transformer的主干网络,如PVT、ViL和Swin Transformer,在部分指标上甚至与参数量大得多的模型表现相当。在基于DETR的端到端检测框架及其变体(Deformable DETR, SMCA-DETR)上,使用Container-Light替代ResNet-50主干网络也带来了大幅度的性能提升。 自监督学习(DINO框架):在DINO自监督学习框架下训练Container-Light和DeiT,使用K近邻分类准确率作为评估指标。在各个训练周期,Container-Light的表示学习效果均优于DeiT,最终达到了71.5%的Top-10 KNN准确率(DeiT为69.6%),再次证明了其在无标签数据上的学习效率。
研究的主要结果是多方面的。在图像分类上,Container和Container-Light在准确率、收敛速度和数据效率上均超越了同规模的Transformer模型,并与经过大量神经架构搜索的先进CNN模型(如EfficientNet)竞争激烈。在目标检测和实例分割上,Container-Light证明了其处理高分辨率输入的有效性和高效性,在不同检测框架和指标上均带来了显著的性能飞跃。在自监督学习上,Container-Light学习到的视觉表征质量更高。这些结果共同支撑了研究的核心结论:通过统一视角设计的、融合静态与动态上下文聚合的Container模块,是一种强大且通用的视觉主干网络构建块。
本研究的主要结论是,成功提供了一个统一的数学框架,将看似迥异的CNN、Transformer和MLP-Mixer架构联系起来,揭示了它们本质上是上下文聚合这一通用操作的不同实例。基于此,提出的Container模块及其网络架构,在科学上为理解和发展视觉基础模型提供了新的视角和工具;在应用上,则提供了一个性能强劲、收敛快速、数据高效且能广泛适用于分类、检测、分割及自监督学习等多种任务的通用视觉主干网络解决方案,具有很高的实用价值。
本研究的亮点在于:第一,理论创新:提出了一个简洁而深刻的统一框架,揭示了主流视觉架构的内在共性,这是对领域认知的重要提升。第二,方法新颖:首创了通过可学习系数混合静态与动态亲和矩阵的Container模块,巧妙地结合了不同架构的优势。第三,性能卓越:由此构建的网络在多个核心视觉任务上实现了当时领先的性能,特别是在保持高效的同时,显著提升了目标检测等下游任务的效果。第四,实用性强:Container-Light的设计针对实际应用中的计算瓶颈,使其能够切实可行地部署到需要高分辨率输入的任务中。第五,框架通用:该框架具备高度的灵活性和可扩展性,能够方便地复现、分析和改进现有架构,为后续研究提供了强大的基础。
其他有价值的内容包括:研究中对学习到的静态亲和矩阵的可视化分析,为神经网络中“归纳偏置的形成”提供了直观证据;通过在DeiT中简单添加静态矩阵(Container-PAM)即可提升性能的实验,进一步验证了静态与动态信息互补的观点;文中还提及了多项同时期致力于融合CNN与Transformer的工作,表明这一研究方向的前沿性和重要性,而本研究提供的统一视角有助于将这些并行工作置于更清晰的上下文之中进行理解。论文最后也坦率地指出了局限性,例如Container在处理极高分辨率输入时仍有困难,而Container-Light前几阶段关闭动态矩阵的设计是手动设定的,未来可探索如何根据任务自动学习此配置。