是時候拋棄 tokenization 了？

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

審核編輯 ：李倩