由于扰动学习可能会遭受「维数灾难」:方差随着扰动维数的增加而增加,并且深度网络中通常有数百万个参数同时发生变化。限制可学习维度数量的一种方法是将网络划分为子模块,每个子模块都有一个单独的损失函数。因此,该研究通过增加局部损失函数的数量来抑制方差,具体包括:
1)Blockwise 损失。首先,该研究将网络深度划分为多个模块。每个模块由几个层组成,在每个模块的末端计算一个损失函数,该损失函数用于更新该模块中的参数。这种方法相当于在模块之间添加了一个「停止梯度(stop gradient)」算子,Belilovsky et al. (2019) 和 Löwe et al. (2019) 曾探索过这种局部贪心损失函数。
2)patchwise 损失。图像等感官输入信号具有空间维度。该研究沿着这些空间维度,逐块应用单独的损失。在 Vision Transformer 架构中(Vaswani et al., 2017; Dosovitskiy et al., 2021),每个空间标记代表图像中的一个 patch。在现代深度网络中,每个空间位置的参数通常是共享的,以提高数据效率并降低内存带宽利用率。虽然简单的权重共享在生物学上是不合理的,但该研究仍然在这项工作中考虑共享权重。通过在 patch 之间添加知识蒸馏 (Hinton et al., 2015) 损失来模拟权重共享的效果是可能的。
3) Groupwise 损失。最后,该研究转向通道(channel)维度。为了创建多个损失,该研究将通道分成多个组,每个组都附加到一个损失函数(Patel et al., 2022)。为了防止组之间相互通信,通道仅连接同一组内的其他通道。
具有局部损失的 LocalMixer 残差块。
特征聚合器
简单地将损失分别应用于空间和通道维度,会带来次优性能,因为每个维度仅包含局部信息。对于分类等标准任务的损失,模型需要输入的全局视图来做出决策。标准架构通过在最终分类层之前执行全局平均池化层,来获得此全局视图。因此,该研究探索了在局部损失函数之前聚合来自其他组和空间块的信息的策略。
特征聚合器设计。
实现
网络架构:该研究提出了更适合局部学习的 LocalMixer 架构。它的灵感来自 MLPMixer (Tolstikhin et al., 2021),它由全连接网络和残差块组成。该研究利用全连接网络,使每个空间块在不干扰其他 patch 的情况下执行计算,这更符合局部学习目标。图 1 显示了高级架构,图 2 显示了一个残差块的详细图。
归一化。在跨不同张量维度的神经网络中执行归一化的方法有很多种(Krizhevsky et al., 2012; Ioffe & Szegedy, 2015; Ba et al., 2016; Ren et al., 2017; Wu & He, 2018)。该研究选择了层归一化的局部变体,它在每个局部空间特征块内进行归一化(Ren et al., 2017)。对于分组的线性层,每组单独进行归一化(Wu & He, 2018)。
该研究通过实验发现这种局部归一化在对比学习中表现更好,并且与监督学习中的层归一化大致相同。局部归一化在生物学上也更合理,因为它不执行全局通信。
通常,归一化层放置在线性层之后。在 MLPMixer(Tolstikhin et al., 2021)中,层归一化被放置在每个残差块的开头。该研究发现最好在每个线性层之前和之后放置归一化,如图 2 所示。实验结果表明,这种设计选择对反向传播没有太大影响,但它允许前向梯度学习更快地学习并实现更低的训练错误。
有效实施复制损失。由于特征聚合和复制损失的设计,组(groups)的简单实现在内存消耗和计算方面可能非常低效。然而,每个空间组实际上计算相同的聚合特征和损失函数。这意味着在执行反向传播和正向梯度时,可以跨损失函数共享大部分计算。该研究实现了自定义 JAX JVP/VJP 函数(Bradbury et al., 2018)并观察到显著的内存节省和复制损失的计算速度提升,否则这在现代硬件上运行是不可行的,结果如下图所示。
复制损失的简单和融合实现时,内存和计算使用情况。
实验
该研究将提出的算法与包括反向传播、反馈对齐和其他前向梯度全局变体在内的替代方案进行比较。反向传播是生物学上难以置信的神谕,因为它计算的是真实梯度,而该方法计算的是噪声梯度。反馈对齐通过使用一组随机向后权重来计算近似梯度。
各项实验的结果如下:
图像分类自监督学习。
有线性 readout 的自监督对比学习。
在不同位置添加局部损失对前向梯度性能的影响。
使用不同数量组时,在 CIFAR-10 上训练时 M/8/* 的错误率。
总结
人们通常认为基于扰动的学习无法扩展到大型深度网络。该研究表明,这在某种程度上是正确的,因为梯度估计方差随着扰动隐藏维度的数量而增长,并且对共享权重扰动而言甚至更遭。
但乐观的是,该研究表明大量的局部贪心损失可以帮助更好地推进梯度学习规模,探索了 blockwise、patchwise 以及 groupwise 以及这三者组合的局部损失,在一个较大的网络中总共有四分之一的损失,表现最好。局部活动扰动前向梯度在更大的网络上比以前的无反向传播算法表现更好。局部损失的想法为不同的损失设计开辟了机会,并阐明了如何在大脑和替代计算设备中寻找生物学上合理的学习算法。