Transformer中FFNs的作用
Transformer 与 FFN
Transformer 的基本单位就是一层 block 这里,一个 block 包含 MSA + FFN,目前公认的说法是,
Attention 作为 token-mixer 做 spatial interaction。 FFN (又称 MLP)在后面作为 channel-mixer 进一步增强 representation。 从 2017 至今,过去绝大部分 Transformer 优化,尤其是针对 NLP tasks 的 Efficient Transformer 都是在 Attention 上的,因为文本有显著的 long sequence 问题。
而 FFN 这里,自从 Transformer 提出基本就是一个 Linear Proj + Activation + Linear Proj 的结构,整体改动十分 incremental。

Activation Function
经历了 ReLU, GeLU,Swish, SwiGLU 等等,基本都是 empirical observations,但都是为了给 representation 加上非线性变换做增强。
ReLU 对 pruning 挺有帮助,尤其是过去对 CNN 做 pruning 的工作,激活值为 0 大致意味着某个 channel 不重要,可以去掉。ReLU 会造成 dead neurons,因此在 Transformer 上逐渐被抛弃。
GeLU 在过去一段时间占比相当大,直到现在 ViT 上使用十分广泛,当然也有用 Swish 的,如 MobileViT。
Gated Linear Units 目前在 LLM 上非常流行,其效果和分析来源于 GLU Variants Improve Transformer。如 PaLM 和 LLaMA 都采用了 SwiGLU, 谷歌的 Gemma 使用 GeGLU。
Linear Projections
为什么 Linear Projection 在 Transformer 里就需要这么大?
常见的说法是 Knowledge Neurons。tokens 在前一层 attention 做 global interaction 之后,通过 FFN 的参数中存放着大量 training 过程中学习到的比较抽象的 knowledge 来进一步 update。
如果 FFN 存储着 Transformer 的 knowledge,那么注定了这个地方不好做压缩加速:
- FFN 变小意味着 model capacity 也变小,大概率会让整体 performance 变得很差。我自己也有过一些 ViT 上的实验 (相信其他人也做过),两个 FC 中间会有个 hidden dimension 的 expansion ratio,一般设置为 4。把这个地方调小会发现怎么都不如大点好。当然太大也不行,因为 FFN 这里的 expansion ratio 决定了整个 Transformer 在推理时的 peak memory consumption,有可能造成 out-of-memory (OOM) error,所以大部分我们看到的 expansion ration 也就在 4 倍,一个比较合适的 performance-memory trade-off.
- FFN 中的 activations 非低秩。过去 convnet 上大家又发现 activations 有明显的低秩特性,所以可以通过 low rank 做加速,如 Kaiming 的这篇文章,如下图所示。但是 FFN 中间的 outputs 很难看出低秩的特性,实际做网络压缩的时候会发现 pruning FFN 的 trade-off 明显不如 convnets,而 unstructured pruning 又对硬件不友好。

FNN 的改动
我们想改动一个 model or module 的时候,无非是两个动机:1)Performance。2)Efficiency。
性能上,目前在 NLP 上可以做 Gated MLP, 如 Mamba 的 block 中,或者 DeepMind 的新结构 Griffin (opens new window)。

但是难说这个地方的性能提升是不是来自于更多的参数量和模型复杂度。
在 CV 上,有个心照不宣的 trick,那就是加 depthwise convolution 引入 locality,试过的朋友都知道这个地方的提升在 CV 任务上有多明显,例如 CIFAR100 上,DeiT-Ti 可以涨接近 10 个点这样子。
但是,鉴于最原始的 FFN 依然是目前采用最广泛的,并且 conv 引入了 inductive bias,破坏了原先 permutation invariant 的 sequence(因为卷积要求规整的 shape,width x height)。大规模 ViT 训练依然没有采用 depthwise conv,如 CLIP, DINOv2, SAM, etc。
效率上,目前最 promising 是改成 Mixture-of-Expert (MoE),但其实 GPT4 和 Mixtral 8x7B 没出来之前基本是 Google 在 solo,没人关注。当然现在时代变了,Mixtral 8x7B 让 MoE 起死回生。
达到 AGI 需要什么结构?
目前这个阶段,没人知道一周以后会有什么大新闻,就像 Sora 悄无声息放出来,一夜之间干掉 U-Net,我也没法说什么结构是最有效的。
总体上,目前没有任何结构能真的完全 beat Transformer,Mamba 目前 也不行,如这篇发现 copy and paste 不太行,scaling 和 in-context 能力也有待查看。
考虑到未来扩展,优秀的结构应该满足这么几个东西,个人按重要性排序:
Scaling Law。如果 model 很难通过 scale up 提升性能,意义不大(针对 AGI 来讲)。但是建议大家不要针对这个地方过度攻击学术界 paper,学术界很难有资源进行这种实验,路都是一步一步踩出来的,提出一个新 architecture 需要勇气和信心,给一些宽容。嗯,说的就是 Mamba。 In-Context Learning 能力。这个能力需要强大的 retrieval 能力和足够的 capacity,而对于 Transformer 来讲,retrieval 靠 Attention,capacity 靠 FFN。scaling 带来的是两者协同提升,进而涌现强大的 in-context learning 能力。 Better Efficiency。说到底这也是为什么我们想换掉 Transformer。做过的朋友都知道 Transformer 训练太耗卡了,无论是 NLP 还是 CV 上。部署的时候又不像 CNN 可以做 bn conv 融合,inference memory 大,low-bit quantization 效果上也不如 CNN,大概率是 attention 这个地方 low-bit 损失大。在满足 1,2 的情况下,如果一个新结构能在 speed, memory 上展现出优势那非常有潜力。Mamba 能火有很大一部分原因是引入 hardware-aware 的实现,极大提升了原先 SSM 的计算效率。 Life-long learning. 知识是不断更新的,训练一个 LLM 需要海量 tokens,强如 OpenAI 也不可能每次 Common Crawl 放出新 data 就从头训一遍,目前比较实际的方案是持续训练,但依然很耗资源。未来的结构需要更高效且持久地学习新知识。
如何理解 Transformers 中 FFNs 的作用?
https://www.zhihu.com/question/622085869/answer/3518358912
原始论文给的答案
FNN 可以看作用 1x1 的卷积核来进行特征的升维和降维。
研究人员发现 FFN 和 ResNet 的 Skip Connection 无论去掉哪一个,模型都会变得不可用,具体可以看《One Wide Feedforward Is All You Need》 这篇论文。所以说 Attention, FFN, ResNet 可以认为是 Transformers 架构的三驾马车,缺一不可。
FFN 设计的初衷,其实就是为模型引入非线性变换。那接着问,attention 中也有 softmax,也是非线性,那 FFN 还是必须的么?
其实这么追问是个陷阱,用这个陷阱其实会引出另外一个问题:attention 是线性运算的还是非线性运算的?
全局来看,对于
但是对于 value 来说,并没有任何的非线性变换。所以每一次 Attention 的计算相当于是对 value 代表的向量进行了加权平均,虽然权重是非线性的权重。这就是 FFN 必须要存在的原因,或者说更本质的原因是因为 FFN 提供了最简单的非线性变换。
Attention, FFN, ResNet 缺一不可但却可能是各司其职,个人的观点(并不一定准确)是, Attention 的功能是做信息的提取和聚合,Resnet 提供信息带宽,而真正学到的知识或者信息都存储在 FFN 中。在图像领域中,也有一种说法,那就是 Attention 其实是 token mixer, FNN 其实是 channel mixer.
功能性的角度
首先大家先猜一下整个 Transformers 中 FFN 的参数占比是多少? 答案是 2/3,你猜对了么?
关于 FFN 的作用,后续研究人员做了很多实验和研究。但是直到目前也不能说就研究清楚了,因为神经网络的解释性本来就差。一个新技术的应用往往要比理论提前好久,比如 GBDT 等 ensemble 模型非常好用,但是很多年以后才用 Margin 的理论研究明白。
这里将一些 FFN 研究的成果做一个小汇总,只列举了目前大家都比较公认的结果。
FFN 是 Transformers 的必备模块,没有 FFN 的 Transformers 学不到什么东西 上面的结论同样适用于 Transformers 中的 skip connect。
《 Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth》这篇论文,提出了 Transformers 架构存在 token uniformity 的归纳偏置(inductive bias,有时候也叫归纳偏好)问题。如果去掉 FFN 或者 Resnet,则问题更加严重。
这里解释一下这两个名词,所谓归纳偏置,可以通俗的理解为模型的“个性”,就是满足训练集合的解法有无数种,但是不同的模型架构会让模型更偏向于某些解法。比如我们常用的一些正则化方法,其实就是让模型的归纳偏置倾向于选择一些简单的解法。任何模型都有归纳偏置,尤其是碰到未见过的样本的时候,模型的归纳偏置就更容易体现出来。 Transformers 的一个归纳偏执是什么呢?就是 token uniformity,有时候也叫 information diffusion,或者 anisotropic (各向异性),也就是说 训练完后的 token 会共享很多相似信息。详细可以看Bert 中的词向量各向异性具体什么意思啊? (opens new window)。
看下图大概就知道了,我们期望表示 token 的向量,相似的要相近,不相似的要远,而且最好是均匀的分布在整个空间中,比如下图所示。但是 Transformers 会存在 各向异性的问题,也就是所有的 token 都挤到一个很窄的锥形区域了。

回到论文,论文将 FFN 和 ResNet 去掉之后做了一些消融实验,证明了 FFN 和 ResNet 是 Transformers 中的必备组件,这两个可以大大的缓解 token uniformity 或者 各向异性的问题。
论文中从数学上证明了经过 Attention 变换后的输出与 rank-1 的矩阵之间的差值存在上界,但是有点复杂,我也没仔细推导过。简单一点的理解呢,就是 Attention 本质上是 value 的线性变换(虽然线性变换的权重是非线性的 softmax)。
当然并不是说 Transformers 就已经将 token uniformity 问题解决了,这个问题依然存在,所以后续又有 Bert-flow 、whitening 等改进。详细可以参考:Bert 中的词向量各向异性具体什么意思啊? - 看图学的回答 - 知乎 https://www.zhihu.com/question/460991118/answer/2353153090
FFN 承担了记忆功能
Transformer Feed-Forward Layers Are Key-Value Memories》
FFN 是一种混合专家模型
MoEfication: Transformer Feed-forward Layers are Mixtures of Experts
这是刘知远团队的论文,其实一直以来,神经网络就存在稀疏激活的现象,也就是在推理的时候,其实只有极小一部分参数参与了计算。这篇论文则通过 MoE 的思想来将 FFN 层拆分成了多个专家,并且新增了一个路由模块来确定推理的时候来挂哪个专家的门诊:)
这么做完之后,在提升推理速度的同时,效果依然能保持原来的 95%以上。