本文目錄

51 無(wú)需訓(xùn)練，Token 合并打造更快的 ViT 架構(gòu)
(來(lái)自佐治亞理工學(xué)院，Meta AI)
51.1 ToMe 論文解讀
51.1.1 背景和動(dòng)機(jī)
51.1.2 Token Merging 的基本思路
51.1.3 什么樣的 tokens 是相似的
51.1.4 Token Merging 的具體步驟：二分軟匹配
51.1.5 Token Merging 的后續(xù)操作：調(diào)節(jié)注意力權(quán)重
51.1.6 其他消融實(shí)驗(yàn)結(jié)果
51.1.7 圖像實(shí)驗(yàn)結(jié)果
51.1.8 視頻實(shí)驗(yàn)結(jié)果

Transformer 是 Google 的團(tuán)隊(duì)在 2017 年提出的一種 NLP 經(jīng)典模型，現(xiàn)在比較火熱的 Bert 也是基于 Transformer。Transformer 模型使用了 Self-Attention 機(jī)制，不采用 RNN 的順序結(jié)構(gòu)，使得模型可以并行化訓(xùn)練，而且能夠擁有全局信息。

51 無(wú)需訓(xùn)練，Token 合并打造更快的 ViT 架構(gòu)

論文名稱(chēng)：Token Merging: Your ViT But Faster

論文地址：https://arxiv.org/pdf/2210.09461.pdf

51.1.1 背景和動(dòng)機(jī)

與卷積神經(jīng)網(wǎng)絡(luò) (CNN) 相比，視覺(jué) Transformer 模型 (ViT) 有一系列優(yōu)良的性質(zhì)，比如：

• Transformer 模型的 Attention 模塊和 MLP 模塊主要有矩陣乘法這種可以加速的操作構(gòu)成。
• Transformer 支持一些性能強(qiáng)大的自監(jiān)督學(xué)習(xí)任務(wù) (掩碼圖像建模 MAE 等等)。
• Transformer 適配多種模態(tài)的輸入數(shù)據(jù) (圖片，文本，音頻等)。
• Transformer 對(duì)于超大規(guī)模數(shù)據(jù)集 (ImageNet-22K) 的泛化性好，預(yù)訓(xùn)練之后的模型在下游任務(wù)中 (比如 ImageNet-1K 圖像分類(lèi)任務(wù)) 表現(xiàn)卓越。

但是在資源受限的邊緣設(shè)備 (如手機(jī)和無(wú)人機(jī)) 上實(shí)際運(yùn)行 Transformer 不太友好，因?yàn)?Transformer 模型又相對(duì)較大的延時(shí)。一種常見(jiàn)的加速視覺(jué) Transformer 模型的方法是對(duì) token (圖片 Patch) 進(jìn)行剪枝。比如 DynamicViT[1]，AdaViT[2]，A-ViT[3]，SPViT[4]等。這些 token Pruning 方法雖然在精度方面能夠?qū)崿F(xiàn)不錯(cuò)的效果，但是 token 剪枝的缺點(diǎn)有：

1. 需要額外的訓(xùn)練過(guò)程，對(duì)資源不友好。
2. token 剪枝限制了模型的實(shí)用性，當(dāng) token 數(shù)量隨著輸入的變化而發(fā)生變化時(shí)，無(wú)法進(jìn)行批處理 (Batch Inference)。為了解決這個(gè)問(wèn)題，大多數(shù) token 剪枝的工作借助了 Mask，對(duì)冗余的 token 進(jìn)行遮擋。但是這樣的做法并沒(méi)有真正剪去這些冗余的 token，使得這些方法并不能在實(shí)際業(yè)務(wù)中真正加速。
3. token 剪枝帶來(lái)的信息損失限制了可以允許剪枝的 token 數(shù)量。

另一種加速 ViT 的做法是對(duì) token (圖片 Patch) 進(jìn)行融合。比如 Token Pooling[5]，Token Learner[6]。和本文方法最接近的 Token Pooling 使用了一個(gè)緩慢的基于 k-means 的方法，但是速度較慢，不適用于現(xiàn)成的模型。

本文希望做一個(gè)無(wú)需訓(xùn)練并且兼顧性能-速度權(quán)衡的 token 融合方法。因?yàn)槠錈o(wú)需訓(xùn)練的優(yōu)良屬性，對(duì)于大模型將會(huì)非常友好。在訓(xùn)練過(guò)程中使用 ToMe，可以觀察到訓(xùn)練速度增長(zhǎng)，總訓(xùn)練時(shí)間縮短了一半。

51.1.2 Token Merging 的基本思路

Token Merging 的基本思路是在一個(gè) ViT 模型中間插入一些 token merging 的模塊, 希望把這些模塊植入 ViT 以后, 訓(xùn)練和推理的速度都有提升。基本作法是在每一個(gè)層之后減少 個(gè) token, 那么一個(gè)有 層的 Transformer 模型從頭到尾減少的 token 數(shù)量就是 。這個(gè) 值越高, 減少的 token 數(shù)量就越多, 但是精度也會(huì)越差。而且值得注意的是, 無(wú)論一張輸入圖片有多少個(gè) tokens, 都會(huì)減少 個(gè) token, 而不是像上文的 token 剪枝算法那樣使得 token 的數(shù)量動(dòng)態(tài)變化。為什么這么設(shè)計(jì)呢? 原因就是如上文所述當(dāng) token 數(shù)量隨著輸入的變化而發(fā)生變化時(shí), 無(wú)法進(jìn)行批處理 (Batch Inference), 使得這些方法并不能在實(shí)際業(yè)務(wù)中真正加速。

如下圖1所示是 Token Merging 的示意圖，ToMe 的位置被插在 Attention 模塊和 MLP 模塊之間，因?yàn)樽髡呦Ｍ柚?Attention 中的特征幫助決定該去融合哪些 tokens。

51.1.3 什么樣的 tokens 是相似的

根據(jù)上面的基本思路，要考慮的第1個(gè)問(wèn)題是我們應(yīng)該合并哪些 tokens，即什么樣的 tokens 可以被認(rèn)為是相似的 tokens？一種比較直接的想法是距離比較近的 tokens 是相似的，但是并不是最優(yōu)解。

如下圖2所示為消融實(shí)驗(yàn)結(jié)果，意在探索什么樣的 tokens 是相似的。消融實(shí)驗(yàn)使用的模型是 MAE 訓(xùn)練策略下得到的 ViT-L/16 預(yù)訓(xùn)練模型 (acc: 85.96%, im/s: 93.3)，不再進(jìn)行任何額外訓(xùn)練。使用合并，這將在網(wǎng)絡(luò)的24層上逐漸移除 98% 的 tokens。

如左圖所示為使用什么特征衡量相似度，作者發(fā)現(xiàn)使用 Key 來(lái)衡量相似度對(duì)性能最友好，因?yàn)?Attention 模塊中的 Key 已經(jīng)總結(jié)了每個(gè) token 中包含的信息，以便用于 Attention 中的 dot-product 相似度。如右圖所示為使用什么距離衡量相似度，作者發(fā)現(xiàn)使用余弦距離來(lái)衡量 token 之間的相似度可以獲得最好的精度-速度權(quán)衡。

如下圖3所示，把不同 head 的 Key 進(jìn)行取平均操作，而不是拼接在一起，更有助于效率。

51.1.4 Token Merging 的具體步驟：二分軟匹配

在定義了 tokens 的相似性之后，下面就需要一種快速的方法來(lái)確定要匹配哪些 tokens，以便在實(shí)際運(yùn)行時(shí)能夠快速將 tokens 的數(shù)量減少 。這個(gè)過(guò)程對(duì)于延時(shí)的要求很高，因?yàn)樵?ViT 模型中要對(duì)可能上千個(gè) tokens 執(zhí)行匹配 次，所以這個(gè)匹配算法的運(yùn)行時(shí)間必須完全可以忽略不計(jì)。

作者在這里為什么不用聚類(lèi)算法 (clustering)，而是匹配算法 (matching) 呢？

答： 因?yàn)榫垲?lèi)算法沒(méi)有限制每個(gè)類(lèi)的具體數(shù)量，因此無(wú)法做到在每層精確地減少 個(gè) tokens。而匹配算法不同，匹配算法可以做到精確地匹配 對(duì) tokens，并把它們?nèi)诤显谝黄?。這樣一來(lái)，就做到了在每層精確地減少 個(gè) tokens。

下面是作者在本文提出的二分軟匹配算法 (Bipartite Soft Matching)，如下圖4所示。

1. 把 ToMe 模塊輸入的所有 tokens 分為相同大小的 2 個(gè)集合 。
2. 把從集合 中的每個(gè) token 到 中與其最相似的 token 畫(huà)一條邊。
3. 只留下最相似的 條邊, 其余刪掉。
4. 融合仍然相連的 條邊 (特征取均值)。
5. 把這兩個(gè)集合拼在一起, 得到 ToMe 模塊的融合結(jié)果。

這個(gè)算法的 PyTorch 代碼如下：

def bipartite_soft_matching(
    metric: torch.Tensor,
    r: int,
    class_token: bool = False,
    distill_token: bool = False,
) -> Tuple[Callable, Callable]:
    """
    Applies ToMe with a balanced matching set (50%, 50%).
    Input size is [batch, tokens, channels].
    r indicates the number of tokens to remove (max 50% of tokens).
    Extra args:
     - class_token: Whether or not there's a class token.
     - distill_token: Whether or not there's also a distillation token.
    When enabled, the class token and distillation tokens won't get merged.
    """
    protected = 0
    if class_token:
        protected += 1
    if distill_token:
        protected += 1

    # We can only reduce by a maximum of 50% tokens
    t = metric.shape[1]
    r = min(r, (t - protected) // 2)

    if r <= 0:
        return do_nothing, do_nothing

    with torch.no_grad():
        metric = metric / metric.norm(dim=-1, keepdim=True)
        a, b = metric[..., ::2, :], metric[..., 1::2, :]
        scores = a @ b.transpose(-1, -2)

        if class_token:
            scores[..., 0, :] = -math.inf
        if distill_token:
            scores[..., :, 0] = -math.inf

        node_max, node_idx = scores.max(dim=-1)
        edge_idx = node_max.argsort(dim=-1, descending=True)[..., None]

        unm_idx = edge_idx[..., r:, :]  # Unmerged Tokens
        src_idx = edge_idx[..., :r, :]  # Merged Tokens
        dst_idx = node_idx[..., None].gather(dim=-2, index=src_idx)

        if class_token:
            # Sort to ensure the class token is at the start
            unm_idx = unm_idx.sort(dim=1)[0]

    def merge(x: torch.Tensor, mode="mean") -> torch.Tensor:
        src, dst = x[..., ::2, :], x[..., 1::2, :]
        n, t1, c = src.shape
        unm = src.gather(dim=-2, index=unm_idx.expand(n, t1 - r, c))
        src = src.gather(dim=-2, index=src_idx.expand(n, r, c))
        dst = dst.scatter_reduce(-2, dst_idx.expand(n, r, c), src, reduce=mode)

        if distill_token:
            return torch.cat([unm[:, :1], dst[:, :1], unm[:, 1:], dst[:, 1:]], dim=1)
        else:
            return torch.cat([unm, dst], dim=1)

    def unmerge(x: torch.Tensor) -> torch.Tensor:
        unm_len = unm_idx.shape[1]
        unm, dst = x[..., :unm_len, :], x[..., unm_len:, :]
        n, _, c = unm.shape

        src = dst.gather(dim=-2, index=dst_idx.expand(n, r, c))

        out = torch.zeros(n, metric.shape[1], c, device=x.device, dtype=x.dtype)

        out[..., 1::2, :] = dst
        out.scatter_(dim=-2, index=(2 * unm_idx).expand(n, unm_len, c), src=unm)
        out.scatter_(dim=-2, index=(2 * src_idx).expand(n, r, c), src=src)

        return out

    return merge, unmerge

其中, 關(guān)鍵變量的含義, 維度和相關(guān)的注釋如下:

src: 集合 A, shape: (B,N,c)

dst: 集合 B, shape: (B,N,c), 其中第1個(gè) token 是 [distillation] token

unm: 集合 A 中不 merge 的 tokens, shape: (B,N?r,c) ，其中第1個(gè) token 是 [class] token

src：集合 A 中要 merge 的 tokens, shape: (B,r,c)

借助 dst.scatter_reduce() 函數(shù)在集合 B 中完成 token merging 操作。

借助 torch.cat() 函數(shù)完成 merge 之后的集合 B 與不 merge 的集合 A 中的 tokens 的拼接工作。

51.1.5 Token Merging 的后續(xù)操作：調(diào)節(jié)注意力權(quán)重

前文提到, ToMe 模塊會(huì)融合 個(gè) token。在 ViT 模型里面, 一個(gè) token 代表輸入圖片的一個(gè) Patch, 比如輸入圖片有 個(gè) Patch, 就是有 個(gè) token。Attention 矩陣的維度也是 的, 它代表了 個(gè) Patch 之間的相關(guān)關(guān)系。但是現(xiàn)在我們?nèi)诤狭? 個(gè) token 之后呢, Attention 矩陣的維度應(yīng)該是 的, 融合了 token 之后, 有的 Key 應(yīng)該占的 Attention 比重大一些, 因?yàn)樗诤狭硕鄠€(gè) token 的信息。

所以作者在這里定義了一個(gè)行向量 是包含每個(gè) token 大小 (token 所代表的 Patch 數(shù)量) 的行向量。

通過(guò)上式將行向量 直接加在 Attention 矩陣上面，相當(dāng)于是人為增加了有些 Key 的 attention weight，而這些 key 恰好是發(fā)生了融合的 Key。

到目前為止，已經(jīng)能夠直接向已經(jīng)訓(xùn)練好的 ViT 模型中添加 ToMe 模塊。使用 ToMe 模塊進(jìn)行訓(xùn)練雖然不是必須的，但是它可以減少準(zhǔn)確度下降，并且加快訓(xùn)練速度。ToMe 模塊本質(zhì)上是 token 的均值操作，因此可以視為是一種池化操作 (Pooling)。因此，我們可以按照平均池化操作 (Average Pooling) 的方式進(jìn)行反向傳播。

51.1.6 其他消融實(shí)驗(yàn)結(jié)果

定義式1所示的對(duì)不同 tokens 進(jìn)行加權(quán)的方式為 weighted avg，在決定合并哪些 tokens 之后，通過(guò)對(duì) tokens 進(jìn)行平均加權(quán)來(lái)合并它們。下圖5左側(cè)的消融實(shí)驗(yàn)結(jié)果表明，weighted avg 的方式優(yōu)于直接的 average pooling 的方式以及 max pooling 的方式。

下圖5右側(cè)的 sequential 代表在將 tokens 分為集合 和集合 時(shí)采用順序的方式，alternating 代表交替分的方式，結(jié)果更好。

如下圖6所示為不同 token 縮減算法的實(shí)驗(yàn)結(jié)果對(duì)比，這些 token 縮減算法包括剪枝，合并等等。Pruning 速度很快，但是隨著 98% 的 token 被移除，重要的信息丟失了。相比之下，合并 token 的算法只有在合并不相似的 token 時(shí)才會(huì)丟失信息。因此，正確地選擇要合并的相似 token 非常重要。

k-means 確實(shí)是一種顯而易見(jiàn)的 token 合并算法，但是除了速度慢的缺點(diǎn)之外，它只比 Pruning 算法好一點(diǎn)。因?yàn)?k-means 算法允許大量的 tokens 被匹配在一起，因此增加了不相似的 tokens 被合并在一起的概率。

相反，本文提出的 merging 算法只會(huì)合并最相似的 tokens，通過(guò)重復(fù) 次逐漸減小 token 的數(shù)量，具有更好的性能和 Pruning 方法的速度優(yōu)勢(shì)。

作者接下來(lái)對(duì)比了15000種 merging 的策略，結(jié)果如下圖7所示。模型是訓(xùn)練好的 AugReg ViT-B/16 模型，以 fp16 在 ImageNet-1K 上的精度。折現(xiàn)是每層融合 個(gè) tokens 的 Constant Merging Schedule。作者發(fā)現(xiàn)，這種均勻融合的策略已經(jīng)接近最優(yōu)了。

51.1.7 圖像實(shí)驗(yàn)結(jié)果

對(duì)于圖像實(shí)驗(yàn)，作者使用 ImageNet-1K 數(shù)據(jù)集，使用了4種不同的方式來(lái)訓(xùn)練，分別是：AugReg[7]，MAE[8]，SWAG[9]，DeiT[10]。

有監(jiān)督和弱監(jiān)督模型實(shí)驗(yàn)結(jié)果

作者在11個(gè) SOTA 的預(yù)訓(xùn)練 ViT 模型 (直接下載開(kāi)源模型，不進(jìn)行任何額外的訓(xùn)練) 上使用了本文提出的 ToMe 方法。AugReg 實(shí)驗(yàn)結(jié)果如下圖8所示，為在大規(guī)模數(shù)據(jù)集預(yù)訓(xùn)練的模型，再在 ImageNet-1K 上 fine-tune 得到的結(jié)果。SWAG 實(shí)驗(yàn)結(jié)果如下圖9所示，為在大規(guī)模數(shù)據(jù)集弱監(jiān)督預(yù)訓(xùn)練的模型，再在 ImageNet-1K 上 fine-tune 得到的結(jié)果。結(jié)果表明，無(wú)論模型的尺寸和類(lèi)型，ToMe 都能夠帶來(lái)約2倍的吞吐量加速。即使減少 96-98% 的 tokens，最大的模型幾乎沒(méi)有任何精度下降：在2倍吞吐量的設(shè)置下，AugReg 得到的 ViT-B，ViT-S 和 ViT-Ti 都有大約 4-5% 的精度下降。ViT-L 在 224px 圖像上僅下降 2%，在 384px 圖像上下降 0.7%，可能是因?yàn)楦蟮妮斎雸D片有更多的 tokens。而且，因?yàn)榇笮湍Ｐ透?，因此允許其中間特征發(fā)生更漸進(jìn)的變化，這也減少了 merging 帶來(lái)的影響。

自監(jiān)督模型實(shí)驗(yàn)結(jié)果

MAE 實(shí)驗(yàn)結(jié)果如下圖10所示，為在大規(guī)模數(shù)據(jù)集弱監(jiān)督預(yù)訓(xùn)練的模型，再在 ImageNet-1K 上 fine-tune 得到的結(jié)果。結(jié)果顯示，在2倍吞吐量的設(shè)置下，MAE 得到的 ViT-H，ViT-L 和 ViT-TB 分別有 0.4%，0.6% 和 1.7% 的精度下降。

與其他模型對(duì)比

如下圖11所示是 ToMe 方法 + MAE 微調(diào)的模型 (具體是在 MAE 進(jìn)行微調(diào)的環(huán)節(jié)用上了本文的 ToMe 方法) 與其他 ImageNet-1K 模型的性能對(duì)比，可以看到 ToMe 方法可以提高 ViT 模型的吞吐量，使得較深的 ViT 模型 (如 ViT-H 和 ViT-L) 的吞吐量與較淺的模型相當(dāng)。

與 Token Pruning 方法對(duì)比

如下圖12所示是 ToMe 方法與 Token Pruning 方法 (DynamicViT[1]，A-ViT[3]，SPViT[4]) 在 DeiT-S 模型上的對(duì)比結(jié)果。ToMe 方法可以再不使用梯度技巧，如 gumbel softmax 等，不添加額外的參數(shù)，以及不使用額外的訓(xùn)練技巧的情況下匹配性能，并超過(guò)現(xiàn)有的 Token Pruning 方法。

而且，Token Pruning 方法通過(guò)由于自身的限制往往使用 token padding 或者 attention 掩碼的方法，使得剪枝帶來(lái)的好處沒(méi)法發(fā)揮出來(lái)。但是，ToMe 方法不受這個(gè)問(wèn)題的影響。

可視化實(shí)驗(yàn)結(jié)果

如下圖13所示是在網(wǎng)絡(luò)的結(jié)尾處的每個(gè)合并的 token 所對(duì)應(yīng)的輸入 Patch?？梢园l(fā)現(xiàn)，ToMe 方法造成的 token 融合的效果和分割很像。比如，在第2張圖中，哈士奇的腿、身體和臉被合并到了不同的 token 中。在第3張圖中，猴子的手、身體、臉、眼睛和嘴都被合并到了不同的 token 中。在最后1張圖中，所有實(shí)例 (狗) 中相同的部分會(huì)被合并在一起。值得注意的是，與剪枝不同，ToMe 這種 token 融合的方法能夠合并背景和前景中的大量冗余的 tokens，而且不丟失信息。

51.1.8 視頻實(shí)驗(yàn)結(jié)果

對(duì)于視頻實(shí)驗(yàn)，作者使用 Kinetics-400 數(shù)據(jù)集，使用了 Spatiotemporal MAE[11] 的方式來(lái)訓(xùn)練。仿照?qǐng)D像實(shí)驗(yàn)的兩種做法進(jìn)行驗(yàn)證，一種是直接把 ToMe 方法應(yīng)用在現(xiàn)成的訓(xùn)練好的模型中，另一種是在 MAE 進(jìn)行微調(diào)的環(huán)節(jié)用上 ToMe 方法。實(shí)驗(yàn)結(jié)果如下圖14所示。將 ToMe 方法應(yīng)用在 ViT-L 上之后，吞吐量與 Swin-B 接近，同時(shí)性能更好。而且，將 ToMe 方法應(yīng)用在 ViT-L 上之后，使用 Spatiotemporal MAE[11] 的方式，性能明顯優(yōu)于 MAE 方式訓(xùn)練的 ViT-B 模型，說(shuō)明 token 融合的方法比 model scaling 更好。

總結(jié)

ToMe 是一個(gè)無(wú)需訓(xùn)練并且兼顧性能-速度權(quán)衡的 token 融合方法，意在縮減 ViT 模型中大量冗余的 tokens。Token Merging 的基本思路是在一個(gè) ViT 模型中間插入一些 token merging 的模塊，希望把這些模塊植入 ViT 以后，訓(xùn)練和推理的速度都有提升。在圖像和視頻中多個(gè)模型的實(shí)驗(yàn)結(jié)果表明，這種 token 融合的方法能夠合并背景和前景中的大量冗余的 tokens，提高 ViT 模型的吞吐量，而且不丟失信息。

參考

1. ^abDynamicViT: Efficient Vision Transformers with Dynamic Token Sparsification
2. ^Adavit: Adaptive vision transformers for efficient image recognition
3. ^abA-ViT: Adaptive tokens for efficient vision transformer
4. ^abSpvit: Enabling faster vision transformers via soft token pruning
5. ^Token pooling in vision transformers
6. ^Tokenlearner: Adaptive space-time tokenization for videos
7. ^How to train your ViT? Data, Augmentation, and Regularization in Vision Transformers
8. ^Masked autoencoders are scalable vision learners
9. ^Revisiting weakly supervised pre-training of visual perception models
10. ^Training data-efficient image transformers & distillation through attention
11. ^abMasked Autoencoders As Spatiotemporal Learners