20 Nov

Transformer升级之路:15、Key归一化助力长度外推

大体上,我们可以将目前Transformer的长度外推技术分为两类:一类是事后修改,比如NTK-RoPEYaRNReRoPE等,这类方法的特点是直接修改推理模型,无需微调就能达到一定的长度外推效果,但缺点是它们都无法保持模型在训练长度内的恒等性;另一类自然是事前修改,如ALIBIKERPLEXPOS以及HWFA等,它们可以不加改动地实现一定的长度外推,但相应的改动需要在训练之前就引入,因此无法不微调地用于现成模型,并且这类方法是否能够Scale Up还没得到广泛认可。

在这篇文章中,笔者将介绍一种意外发现的长度外推方案——“KeyNorm”——对Attention的Key序列做L2 Normalization,很明显它属于事前修改一类,但对Attention机制的修改非常小,因此看上去非常有希望能够Scale Up。

最初动机

之所以说“意外发现”,是因为该改动的原始动机并不是长度外推,而是尝试替换Scaled Dot-Product Attention中的Scale方式。我们知道,Attention的标准定义是(本文主要考虑Causal场景)
\begin{equation}\boldsymbol{o}_i = \frac{\sum_{j = 1}^i\exp\left(\frac{\boldsymbol{q}_i\cdot \boldsymbol{k}_j}{\sqrt{d}}\right)\boldsymbol{v}_j}{\sum_{j = 1}^i\exp\left(\frac{\boldsymbol{q}_i\cdot \boldsymbol{k}_j}{\sqrt{d}}\right)},\quad \boldsymbol{q}_i,\boldsymbol{k}_j\in\mathbb{R}^d\label{eq:sdpa}\end{equation}

点击阅读全文...

9 Nov

VQ一下Key,Transformer的复杂度就变成线性了

Efficient Transformer,泛指一切致力于降低Transformer的二次复杂度的工作,开始特指针对Attention的改进,后来更一般的思路,如傅立叶变换、线性RNN等,也被归入这个范畴。不得不说,为了降低Transformer的二次复杂度,各路大牛可谓是“八仙过海,各显神通”,各种神奇的思路“百花齐放”,笔者也从中学习到了不少理论知识。然而,尽管Efficient Transformer在理论上是精彩的,但实际上该领域一直都是不愠不火的状态,并没有实际表现十分出色的模型,在LLM火爆的今天,甚至已经逐渐淡出了大家的视野,也淡出了笔者的兴趣范围。

不过,最近有一篇论文《Transformer-VQ: Linear-Time Transformers via Vector Quantization》,却让笔者为之拍案叫绝。作者非常高明地洞察到,只需要对标准Attention的Key做一下VQ(Vector Quantize),复杂度就会自动降低为线性!这种线性化思路保留了标准Attention的形式,是标准Attention到线性Attention的一个完美过渡,同时最大程度上保留了标准Attention的能力。

高效难题

说起来,本站也算是比较早关注Efficient Transformer相关工作了,最早可以追溯到2019年解读Sparse Transformer的一篇博客《为节约而生:从标准Attention到稀疏Attention》。此后,陆续写的关于Efficient Transformer的其他博文还有

点击阅读全文...

31 Oct

简单得令人尴尬的FSQ:“四舍五入”超越了VQ-VAE

正如“XXX is all you need”一样,有不少论文都以“简单得令人尴尬”命名(An Embarrassingly Simple XXX),但在笔者看来,这些论文大多数都是噱头多于实力。不过,笔者最近阅读到的一篇论文,真的让人不由得发出“简单得令人尴尬”的感叹~

论文的标题是《Finite Scalar Quantization: VQ-VAE Made Simple》,顾名思义,这是一篇旨在用FSQ(Finite Scalar Quantization)简化VQ-VAE的工作。随着生成模型、多模态LLM的逐渐流行,VQ-VAE及其后续工作也作为“图像的Tokenizer”而“水涨船高”。然而,VQ-VAE的训练本身也存在一些问题,而FSQ这篇论文则声称通过更简单的“四舍五入”就可以达到同样的目的,并且有着效果更好、收敛更快、训练更稳的优点。

FSQ真有这么神奇?接下来我们一起学习一下。

VQ

首先,我们来了解一下“VQ”。VQ全称是“Vector Quantize”,可以翻译为“向量量子化”或者“向量量化”,是指将无限、连续的编码向量映射为有限、离散的整数数字的一种技术。如果我们将VQ应用在自编码器的中间层,那么可以在压缩输入大小的同时,让编码结果成为一个离散的整数序列。

点击阅读全文...

22 Oct

从梯度最大化看Attention的Scale操作

我们知道,Scaled Dot-Product Attention的Scale因子是$\frac{1}{\sqrt{d}}$,其中$d$是$\boldsymbol{q},\boldsymbol{k}$的维度。这个Scale因子的一般解释是:如果不除以$\sqrt{d}$,那么初始的Attention就会很接近one hot分布,这会造成梯度消失,导致模型训练不起来。然而,可以证明的是,当Scale等于0时同样也会有梯度消失问题,这也就是说Scale太大太小都不行。

那么多大的Scale才适合呢?$\frac{1}{\sqrt{d}}$是最佳的Scale了吗?本文试图从梯度角度来回答这个问题。

已有结果

《浅谈Transformer的初始化、参数化与标准化》中,我们已经推导过标准的Scale因子$\frac{1}{\sqrt{d}}$,推导的思路很简单,假设初始阶段$\boldsymbol{q},\boldsymbol{k}\in\mathbb{R}^d$都采样自“均值为0、方差为1”的分布,那么可以算得
\begin{equation}\mathbb{V}ar[\boldsymbol{q}\cdot\boldsymbol{k}] = d\end{equation}

点击阅读全文...

16 Oct

随机分词再探:从Viterbi Sampling到完美采样算法

在文章《随机分词浅探:从Viterbi Decoding到Viterbi Sampling》中,笔者提出了一种名为“Viterbi Sampling”的随机分词算法,它只是在求最优解的Viterbi Decoding基础上进行小修改,保留了Viterbi算法的简单快速的特点,相比于已有的Subword Regularization明显更加高效。不过,知乎上的读者 @鶴舞 指出,当前的采样算法可能会在多次二选一“稀释”了部分方案的出现概率,直接后果是原本分数最高的切分并不是以最高概率出现。

经过仔细思考后,笔者发现相应的问题确实存在,当时为了尽快得到一种新的采样算法,在细节上的思考和处理确实比较粗糙。为此,本文将进一步完善Viterbi Sampling算法,并证明完善后的算法在效果上可以跟Subword Regularization等价的。

问题分析

首先,我们来看一下评论原话

点击阅读全文...

13 Oct

EMO:基于最优传输思想设计的分类损失函数

众所周知,分类任务的标准损失是交叉熵(Cross Entropy,等价于最大似然MLE,即Maximum Likelihood Estimation),它有着简单高效的特点,但在某些场景下也暴露出一些问题,如偏离评价指标、过度自信等,相应的改进工作也有很多,此前我们也介绍过一些,比如《再谈类别不平衡问题:调节权重与魔改Loss的对比联系》《如何训练你的准确率?》《缓解交叉熵过度自信的一个简明方案》等。由于LLM的训练也可以理解为逐token的分类任务,默认损失也是交叉熵,因此这些改进工作在LLM流行的今天依然有一定的价值。

在这篇文章中,我们介绍一篇名为《EMO: Earth Mover Distance Optimization for Auto-Regressive Language Modeling》的工作,它基于最优传输思想提出了新的改进损失函数EMO,声称能大幅提高LLM的微调效果。其中细节如何?让我们一探究竟。

点击阅读全文...

8 Oct

预训练一下,Transformer的长序列成绩还能涨不少!

作为LLM的主流模型架构,Transformer在各类任务上的总体表现都出色,大多数情况下,Transformer的槽点只是它的平方复杂度,而不是效果——除了一个名为Long Range Arena(下面简称LRA)的Benchmark。一直以来,LRA一直是线性RNN类模型的“主场”,与之相比Transformer在上面有明显的差距,以至于让人怀疑这是否就是Transformer的固有缺陷。

不过,近日论文《Never Train from Scratch: Fair Comparison of Long-Sequence Models Requires Data-Driven Priors》将这“缺失的一环”给补齐了。论文指出,缺乏预训练是Transformer在LRA上效果较差的主要原因,而所有架构都可以通过预训练获得一定的提升,Transformer的提升则更为明显。

旧背景

Long Range Arena(LRA)是长序列建模的一个Benchmark,提出自论文《Long Range Arena: A Benchmark for Efficient Transformers》,从论文标题就可以看出,LRA是为了测试各种Efficient版的Transformer而构建的,里边包含了多种类型的数据,序列长度从1k到16k不等,此前不少Efficient Transformer的工作也都在LRA进行了测试。虽然在代表性方面有些争议,但LRA依然不失为一个测试Efficient Transformer的长序列能力的经典Benchmark。

点击阅读全文...

26 Sep

脑洞大开:非线性RNN居然也可以并行计算?

近年来,线性RNN由于其可并行训练以及常数推理成本等特性,吸引了一定研究人员的关注(例如笔者之前写的《Google新作试图“复活”RNN:RNN能否再次辉煌?》),这让RNN在Transformer遍地开花的潮流中仍有“一席之地”。然而,目前看来这“一席之地”只属于线性RNN,因为非线性RNN无法高效地并行训练,所以在架构之争中是“心有余而力不足”。

不过,一篇名为《Parallelizing Non-Linear Sequential Models over the Sequence Length》的论文有不同的看法,它提出了一种迭代算法,宣传可以实现非线性RNN的并行训练!真有如此神奇?接下来我们一探究竟。

求不动点

原论文对其方法做了非常一般的介绍,而且其侧重点是PDE和ODE,这里我们直接从RNN入手。考虑常见的简单非线性RNN:
\begin{equation}x_t = \tanh(Ax_{t-1} + u_t)\label{eq:rnn}\end{equation}

点击阅读全文...