是時(shí)候拋棄 tokenization 了？

ChatGPT 等對(duì)話 AI 的出現(xiàn)讓人們習(xí)慣了這樣一件事情：輸入一段文本、代碼或一張圖片，對(duì)話機(jī)器人就能給出你想要的答案。但在這種簡(jiǎn)單的交互方式背后，AI 模型要進(jìn)行非常復(fù)雜的數(shù)據(jù)處理和運(yùn)算，tokenization 就是比較常見的一種。 在自然語言處理領(lǐng)域，tokenization 指的是將文本輸入分割成更小的單元，稱為「token」。這些 token 可以是詞、子詞或字符，取決于具體的分詞策略和任務(wù)需求。例如，如果對(duì)句子「我喜歡吃蘋果」執(zhí)行 tokenization 操作，我們將得到一串 token 序列：["我", "喜歡", "吃", "蘋果"]。有人將 tokenization 翻譯成「分詞」，但也有人認(rèn)為這種翻譯會(huì)引起誤導(dǎo)，畢竟分割后的 token 未必是我們?nèi)粘Ｋ斫獾摹冈~」。

圖源：https://towardsdatascience.com/dynamic-word-tokenization-with-regex-tokenizer-801ae839d1cd Tokenization 的目的是將輸入數(shù)據(jù)轉(zhuǎn)換成計(jì)算機(jī)可以處理的形式，并為后續(xù)的模型訓(xùn)練和分析提供一種結(jié)構(gòu)化的表示方式。這種方式為深度學(xué)習(xí)研究帶來了便利，但同時(shí)也帶來了很多麻煩。前段時(shí)間剛加入 OpenAI 的 Andrej Karpathy 指出了其中幾種。 首先，Karpathy 認(rèn)為，Tokenization 引入了復(fù)雜性：通過使用 tokenization，語言模型并不是完全的端到端模型。它需要一個(gè)獨(dú)立的階段進(jìn)行 tokenization，該階段有自己的訓(xùn)練和推理過程，并需要額外的庫。這增加了引入其他模態(tài)數(shù)據(jù)的復(fù)雜性。

此外，tokenization 還會(huì)在某些場(chǎng)景下讓模型變得很容易出錯(cuò)，比如在使用文本補(bǔ)全 API 時(shí)，如果你的 prompt 以空格結(jié)尾，你得到的結(jié)果可能大相徑庭。

圖源：https://blog.scottlogic.com/2021/08/31/a-primer-on-the-openai-api-1.html 再比如，因?yàn)?tokenization 的存在，強(qiáng)大的 ChatGPT 竟然不會(huì)將單詞反過來寫（以下測(cè)試結(jié)果來自 GPT 3.5）。

這樣的例子可能還有很多。Karpathy 認(rèn)為，要解決這些問題，我們首先要拋棄 tokenization。 Meta AI 發(fā)表的一篇新論文探討了這個(gè)問題。具體來說，他們提出了一種名為「 MEGABYTE」的多尺度解碼器架構(gòu)，可以對(duì)超過一百萬字節(jié)的序列進(jìn)行端到端可微建模。

論文鏈接：https://arxiv.org/pdf/2305.07185.pdf 重要的是，該論文展現(xiàn)出了拋棄 tokenization 的可行性，被 Karpathy 評(píng)價(jià)為「很有前途（Promising）」。 以下是論文的詳細(xì)信息。 論文概覽 在 機(jī)器學(xué)習(xí)的文章 中講過，機(jī)器學(xué)習(xí)之所以看上去可以解決很多復(fù)雜的問題，是因?yàn)樗堰@些問題都轉(zhuǎn)化為了數(shù)學(xué)問題。

而 NLP 也是相同的思路，文本都是一些「非結(jié)構(gòu)化數(shù)據(jù)」，我們需要先將這些數(shù)據(jù)轉(zhuǎn)化為「結(jié)構(gòu)化數(shù)據(jù)」，結(jié)構(gòu)化數(shù)據(jù)就可以轉(zhuǎn)化為數(shù)學(xué)問題了，而分詞就是轉(zhuǎn)化的第一步。 由于自注意力機(jī)制和大型前饋網(wǎng)絡(luò)的成本都比較高，大型 transformer 解碼器 (LLM) 通常只使用數(shù)千個(gè)上下文 token。這嚴(yán)重限制了可以應(yīng)用 LLM 的任務(wù)集。 基于此，來自 Meta AI 的研究者提出了一種對(duì)長(zhǎng)字節(jié)序列進(jìn)行建模的新方法 ——MEGABYTE。該方法將字節(jié)序列分割成固定大小的 patch，和 token 類似。 MEGABYTE 模型由三部分組成：

1. patch 嵌入器，它通過無損地連接每個(gè)字節(jié)的嵌入來簡(jiǎn)單地編碼 patch；
2. 全局模塊 —— 帶有輸入和輸出 patch 表征的大型自回歸 transformer；
3. 局部模塊 —— 一個(gè)小型自回歸模型，可預(yù)測(cè) patch 中的字節(jié)。

至關(guān)重要的是，該研究發(fā)現(xiàn)對(duì)許多任務(wù)來說，大多數(shù)字節(jié)都相對(duì)容易預(yù)測(cè)（例如，完成給定前幾個(gè)字符的單詞），這意味著沒有必要對(duì)每個(gè)字節(jié)都使用大型神經(jīng)網(wǎng)絡(luò)，而是可以使用小得多的模型進(jìn)行 intra-patch 建模。

MEGABYTE 架構(gòu)對(duì)長(zhǎng)序列建模的 Transformer 進(jìn)行了三項(xiàng)主要改進(jìn)： 1. sub-quadratic 自注意力。大多數(shù)關(guān)于長(zhǎng)序列模型的工作都集中在減少自注意力的二次成本上。通過將長(zhǎng)序列分解為兩個(gè)較短的序列和最佳 patch 大小，MEGABYTE 將自注意力機(jī)制的成本降低到，即使是長(zhǎng)序列也能易于處理。

2. per-patch 前饋層。在 GPT-3 等超大模型中，超過 98% 的 FLOPS 用于計(jì)算 position-wise 前饋層。MEGABYTE 通過給 per-patch（而不是 per-position）使用大型前饋層，在相同的成本下實(shí)現(xiàn)了更大、更具表現(xiàn)力的模型。在 patch 大小為 P 的情況下，基線 transformer 將使用具有 m 個(gè)參數(shù)的相同前饋層 P 次，而 MEGABYTE 僅需以相同的成本使用具有 mP 個(gè)參數(shù)的層一次。 3. 并行解碼。transformer 必須在生成期間串行執(zhí)行所有計(jì)算，因?yàn)槊總€(gè)時(shí)間步的輸入是前一個(gè)時(shí)間步的輸出。通過并行生成 patch 的表征，MEGABYTE 在生成過程中實(shí)現(xiàn)了更大的并行性。例如，具有 1.5B 參數(shù)的 MEGABYTE 模型生成序列的速度比標(biāo)準(zhǔn)的 350M 參數(shù) transformer 快 40%，同時(shí)在使用相同的計(jì)算進(jìn)行訓(xùn)練時(shí)還改善了困惑度（perplexity）。 總的來說，MEGABYTE 讓我們能夠以相同的計(jì)算預(yù)算訓(xùn)練更大、性能更好的模型，將能夠處理非常長(zhǎng)的序列，并提高部署期間的生成速度。 MEGABYTE 還與現(xiàn)有的自回歸模型形成鮮明對(duì)比，后者通常使用某種形式的 tokenization，其中字節(jié)序列被映射成更大的離散 token（Sennrich et al., 2015; Ramesh et al., 2021; Hsu et al., 2021） 。tokenization 使預(yù)處理、多模態(tài)建模和遷移到新領(lǐng)域變得復(fù)雜，同時(shí)隱藏了模型中有用的結(jié)構(gòu)。這意味著大多數(shù) SOTA 模型并不是真正的端到端模型。最廣泛使用的 tokenization 方法需要使用特定于語言的啟發(fā)式方法（Radford et al., 2019）或丟失信息（Ramesh et al., 2021）。因此，用高效和高性能的字節(jié)模型代替 tokenization 將具有許多優(yōu)勢(shì)。 該研究對(duì) MEGABYTE 和一些強(qiáng)大的基線模型進(jìn)行了實(shí)驗(yàn)。實(shí)驗(yàn)結(jié)果表明，MEGABYTE 在長(zhǎng)上下文語言建模上的性能可與子詞模型媲美，并在 ImageNet 上實(shí)現(xiàn)了 SOTA 的密度估計(jì)困惑度，并允許從原始音頻文件進(jìn)行音頻建模。這些實(shí)驗(yàn)結(jié)果證明了大規(guī)模無 tokenization 自回歸序列建模的可行性。 MEGABYTE 主要組成部分

patch 嵌入器 大小為 P 的 patch 嵌入器能夠?qū)⒆止?jié)序列映射成一個(gè)長(zhǎng)度為、維度為的 patch 嵌入序列。

首先，每個(gè)字節(jié)都嵌入了一個(gè)查找表，形成一個(gè)大小為 D_G 的嵌入，并添加了位置嵌入。

然后，字節(jié)嵌入被重塑成維度為的 K 個(gè) patch 嵌入的序列。為了允許自回歸建模，該 patch 序列被填充以從可訓(xùn)練的 patch 大小的填充嵌入（），然后從輸入中移除最后一個(gè) patch。該序列是全局模型的輸入，表示為。

全局模塊 全局模塊是一個(gè) decoder-only 架構(gòu)的 P?D_G 維 transformer 模型，它在 k 個(gè) patch 序列上進(jìn)行操作。全局模塊結(jié)合自注意力機(jī)制和因果掩碼來捕獲 patch 之間的依賴性。全局模塊輸入 k 個(gè) patch 序列的表示，并通過對(duì)先前 patch 執(zhí)行自注意力來輸出更新的表示。

最終全局模塊的輸出包含 P?D_G 維的 K 個(gè) patch 表示。對(duì)于其中的每一個(gè)，研究者將它們重塑維長(zhǎng)度為 P、維度為 D_G 的序列，其中位置 p 使用維度 p?D_G to (p + 1)?D_G。然后將每個(gè)位置映射到具有矩陣的局部模塊維度，其中 D_L 為局部模塊維度。接著將這些與大小為 D_L 的字節(jié)嵌入相結(jié)合，用于下一個(gè)的 token。

局部字節(jié)嵌入通過可訓(xùn)練的局部填充嵌入（E^local-pad ∈ R^DL）偏移 1，從而允許在 path 中進(jìn)行自回歸建模。最終得到張量 ?

局部模塊 局部模塊是一個(gè)較小的、decoder-only 架構(gòu)的 D_L 維 transformer 模型，它在包含 P 個(gè)元素的單個(gè) patch k 上運(yùn)行，每個(gè)元素又是一個(gè)全局模塊輸出和序列中前一個(gè)字節(jié)的嵌入的總和。K 個(gè)局部模塊副本在每個(gè) patch 上獨(dú)立運(yùn)行，并在訓(xùn)練時(shí)并行運(yùn)行，從而計(jì)算表示。

最后，研究者可以計(jì)算每個(gè)位置的詞匯概率分布。第 k 個(gè) patch 的第 p 個(gè)元素對(duì)應(yīng)于完整序列的元素 t，其中 t = k?P + p。

效率分析 訓(xùn)練效率 在縮放序列長(zhǎng)度和模型大小時(shí)，研究者分析了不同架構(gòu)的成本。如下圖 3 所示，MEGABYTE 架構(gòu)在各種模型大小和序列長(zhǎng)度上使用的 FLOPS 少于同等大小的 transformer 和線性 transformer，允許相同的計(jì)算成本下使用更大的模型。

生成效率 考慮這樣一個(gè) MEGABYTE 模型，它在全局模型中有 L_global 層，在局部模塊中有 L_local 層，patch 大小為 P，并與具有 L_local + L_global 層的 transformer 架構(gòu)進(jìn)行比較。用 MEGABYTE 生成每個(gè) patch 需要一個(gè) O (L_global + P?L_local) 串行操作序列。當(dāng) L_global ≥ L_local（全局模塊的層多于局部模塊）時(shí)，MEGABYTE 可以將推理成本降低近 P 倍。 實(shí)驗(yàn)結(jié)果 語言建模 研究者在強(qiáng)調(diào)長(zhǎng)程依賴的 5 個(gè)不同數(shù)據(jù)集上分別評(píng)估了 MEGABYTE 的語言建模功能，它們是 Project Gutenberg (PG-19)、Books、Stories、arXiv 和 Code。結(jié)果如下表 7 所示，MEGABYTE 在所有數(shù)據(jù)集上的表現(xiàn)始終優(yōu)于基線 transformer 和 PerceiverAR 。

研究者還擴(kuò)展了在 PG-19 上的訓(xùn)練數(shù)據(jù)，結(jié)果如下表 8 所示，MEGABYTE 顯著優(yōu)于其他字節(jié)模型，并可與子詞（subword）上訓(xùn)練的 SOTA 模型相媲美。

圖像建模 研究者在 ImageNet 64x64 數(shù)據(jù)集上訓(xùn)練了一個(gè)大型 MEGABYTE 模型，其中全局和局部模塊的參數(shù)分別為 2.7B 和 350M，并有 1.4T token。他們估計(jì)，訓(xùn)練該模型所用時(shí)間少于「Hawthorne et al., 2022」論文中復(fù)現(xiàn)最佳 PerceiverAR 模型所需 GPU 小時(shí)數(shù)的一半。如上表 8 所示，MEGABYTE 與 PerceiverAR 的 SOTA 性能相當(dāng)?shù)耐瑫r(shí)，僅用了后者一半的計(jì)算量。 研究者比較了三種 transformer 變體，即 vanilla、PerceiverAR 和 MEGABYTE，以測(cè)試在越來越大圖像分辨率上長(zhǎng)序列的可擴(kuò)展性。結(jié)果如下表 5 所示，在這一計(jì)算控制設(shè)置下，MEGABYTE 在所有分辨率上均優(yōu)于基線模型。

下表 14 總結(jié)了每個(gè)基線模型使用的精確設(shè)置，包括上下文長(zhǎng)度和 latent 數(shù)量。

音頻建模 音頻兼具文本的序列結(jié)構(gòu)和圖像的連續(xù)屬性，這對(duì) MEGABYTE 而言是一個(gè)有趣的應(yīng)用。本文模型獲得 3.477 的 bpb，明顯低于 perceiverAR（3.543）和 vanilla transformer 模型（3.567）。更多消融結(jié)果詳見下表 10。

審核編輯 ：李倩