基于 Transformer 的大语言模型，有个很重要的指标就是上下文长度。这是因为模型本身是无状态的，它在推理时只能“记住”上下文窗口里的内容。因此我们在和模型对话时，需要把对话的历史（上下文）和我们的提示一起发给模型，否则模型就只能看到提示本身。随着对话的进行，这个上下文会越来越长，当超过模型的上下文窗口长度时，模型就只能“忘记”最早的对话内容了。

例如在下面这个会话里，如果我们把上下文窗口的长度设为 10 个 token，可以看到模型就不能记住之前对话的内容了：

User> Hi, my name is Ping
AI> Hello Ping, it's nice to meet you! My name is OpenAI.
    How can I assist you today?

User> What is 1+1?
AI> The answer to 1+1 is 2.

User> What is my name?
AI> I'm sorry, I don't have access to that information.
    Could you please tell me your name?

上面这个例子只是文本对话，对于多模态大模型而言（例如视觉大模型），输入的序列长度可以达到几百万个 token，因此上下文长度对于多模态大模型有极其重要的意义。但是，目前基于 Transformer 的大模型，普遍都只支持几千个 token 的上下文长度。

这是因为，Transformer 的计算成本，是和上下文长度高度相关的。每个 Transformer 模块，都包含 Self-Attention 和 Feed Forward Network (FFN) 两个主要的计算部分：

(图源：https://jalammar.github.io/illustrated-transformer/)

其中：

• Self-attention 部分，计算量和上下文长度的平方成正比
• FFN 部分，计算量为 2mT，m 为参数数量，T为序列长度（上下文长度）

大模型的参数量巨大，要支持的序列越长，其训练和推理成本也就越高，因此很难 scale 到更长的上下文，比如几百万 token 的长度。

除此之外，Transformer 大模型还有一个关键依赖，就是分词器（tokenizer）。

大语言模型在工作时，先要把输入文字分成一个个 token，这里的 token 大致相当于单词（有些分词器也会分成 subword），所有见过的 token 集合就是词汇表 (vocabulary)，每个 token 到这个词汇表里查到自己的 id，然后通过 embedding 模型转成一个向量，因此输入序列就是一串向量，每个向量对应一个 token。这里就带来几个问题：

• 这个分词器和大模型是分离的，例如你如果到 HuggingFace 上看，很多模型在用的时候需要指定一个外部的 tokenizer。
• 每个分词器的实现不同，不同的分词器对模型的性能会有影响。
• 由于自然语言本身的模糊性，一段话可能有不同分词方法，有时候分词器需要有领域知识（domain knowledge）才能正确分词。
• 分词器的存在，使得大模型的训练和推理并不是真正端到端（End-to-End）的，这使得大模型的开发和部署都变得更加复杂。

解决这些问题，需要在模型的算法和架构上进行创新，Meta 提出的 MegaByte 就是一个最新的尝试。

讨论一下 MegaByte 的计算复杂性。

注意力部分，MegaByte 有 Global 和 Local 两级。

• Global Model 的序列长度是 K（K 个 patch），因此其计算复杂度为 \(O(K^2)\) ，假设输入序列长度（上下文窗口）为 T，每个 patch 大小为 P，那么 K=T/P，因此 Global Model 的计算复杂度为 \(O(\frac{T^2}{P^2})\) 。
• Local Model 一共有 K 个副本，每个的序列长度为 P（patch 大小），因此计算复杂性为 \(O(KP^2) = O(\frac{T}{P}P^2)=O(TP)\) 。

两者加起来，注意力的计算量为 \(O(\frac{T^2}{P^2}+TP)\) 。这里的 P 是一个用户选择的超参数，当 \(P=T^{1/3}\) 的时候，注意力的计算量为 \(O(T^{3/4})\) ，相比经典 Transformer 的 \(O(T^2)\) 复杂度是一个显著的提升。

前馈网络（FFN）部分：实际上大模型计算量主要是在 FFN，例如 GPT-3 规模的大模型，自注意力只占了 FLOPS 的 1.4%。根据 OpenAI 的论文『Scaling laws for neural language models』，前馈网络的 FLOPS 大致可以估算为 \(2mT\) ，其中 m 是模型的非 embedding 参数量，T是序列长度。

MegaByte 有 Global Model 和 Local Model，假设它们的参数量分别为 \(m_g\) 和 \(m_l\) ，那么 FFN 的计算量分别是：

• Global Model: \(2\frac{T}{P} m_g\) ，P为 patch 大小
• Local Model: 每个序列长度为 P，一共有 K 个（K=T/P），所以计算量为 \(2P m_l \times \frac{T}{P} = 2Tm_l\)

加起来总计算量为 \(2T(\frac{m_g}{P} + m_l)\) 。假设 Local Model 远小于 Global Model，即 \(m_l \ll m_g\) ，那么总计算量可以近似为 \(2T\frac{m_g}{P}\) ，也就是经典 Transformer 的 1/P。

这意味着用同样的算力（FLOPS）预算，MegaByte 可以支持大的多的模型，或者反过来，同样的参数量，MegaByte 的算力需求要低得多。

在生成（推理）方面，作者比较了生成每个 patch（近似于 token）的计算量。

• MegaByte: 生成每个 patch 需要经过 \(O(L_{global} + P \cdot L_{local})\) 个操作， \(L_{global}, L_{local}\) 分别是 Global Model 和 Local Model 的层数。
• 经典 Transformer，假设总层数相同，也有 \(L_{global} + L_{local}\) 层，那么生成 1 个 patch 需要经过的操作数为 \(O(P \cdot L_{global} + P \cdot L_{local})\) 。

同样假设 Global Model 层数比 Local Model 大的多，MegaByte 生成一个 patch 经过的操作可近似为 \(O(L_{global})\) ，是经典 Transformer 的 1/P。

注：这里的计算我感觉不太对，经典 Transformer 的序列长度是 token 为单位的，但是论文里的比较都用 T（字节数）作为序列长度来估算 Transformer 的计算复杂度。当然 Transformer 的 token 一般都比较短（word 或者 subword），如果 token 长度远小于 patch 长度的话，那么 MegaByte 在 FFN 和生成上确实能节省不少 FLOPS，但应该没有 P 倍这么多。

除了 MegaByte 的基本架构，因此作者也讨论了一些变体和优化措施。

MegaByte 的设计目标之一就是多模态，所以作者除了语言模型，还对 MegaByte 在图像和音频上的模型进行了测试，有兴趣的话可以看一下。

最后看一下论文里附带的伪代码。作者只给出了 forward 部分，而且没有注释，所以我只能根据我的理解猜一下，加上了注释。

class MegaByteDecoder:
    def __init__(
            self,
            global_args,
            local_args,
            patch_size,
    ):
        self.pad = 0
        self.patch_size = patch_size
        self.globalmodel = TransformerDecoder(global_args)
        self.localmodel = TransformerDecoder(local_args)

    def forward(
            self,
            bytes, ):
        bytes_global, bytes_local = self.prepare_input(bytes)
        global_bytes_embedded = self.globalmodel.embed(bytes_global)
        # b: batch size
        # t: 序列长度（patch数, 应该也是local model副本数）
        # p: patch长度
        # e: 每个字节embedding长度
        global_in = rearrange(
            global_bytes_embedded,
            # 输入形状：
            #   b: batch里有b个序列，
            #   (t p): 每个序列有t*p个字节
            #   e: 每个字节embeding长度为e
            # 变换为:
            #   b: batch里有b个序列，
            #   t: 每个序列有t个patch
            #   (p e): 每个patch长度为p*e
            "b (t p) e -> b t (p e)",
            p=self.patch_size,
        )
        global_output = self.globalmodel(global_in)

        global_output_reshaped = rearrange(
            global_output,
            # 将输出变成(b*t)个patch，每个patch有p字节，
            # 每个字节embedding大小e
            "b t (p e) -> (b t) p e",
            p=self.patch_size,
        )
        local_bytes_embedded = self.localmodel.embed(bytes_local)
        local_in = local_bytes_embedded + global_output_reshaped
        local_output = self.localmodel(local_in)

        batch_size = bytes_global.shape[0]
        # local的输出是(b*t)个patch，l应该就是前面的p, v应该就是前面的e？
        # reshape成b个序列，每个序列有(t*l)个字节embedding
        x = rearrange(local_output, "(b t) l v  -> b (t l) v", b=batch_size)
        return x

    # 这个函数将字节序列转成global和local能用的形状，注意这里是字节序列，
    # 而不是embedding。所以bytes的形状是(b,T)的2维数组，b是batch size，
    # T是序列长度（字节数）
    # 返回两个东西：
    # - bytes_global是bytes同样形状，但是扔掉了最后一个patch，前面插了padding
    # - bytes_local是(b*t,p)形状，因为这个是给local model用的，每个p字节
    def prepare_input(self, bytes):
        # 看后面的操作，bytes应该是(b,T)的2维数组，b为batch大小，T为序列的长度（字节数）
        # padding_global形状，应该是(b,p)
        padding_global = bytes.new(bytes.shape[0], self.patch_size).fill_(self.pad)

        # 这里就是去掉最后一个patch，前面插上padding
        bytes_global = torch.cat((padding_global, bytes[:, : -self.patch_size]), -1)

        # bytes_input形状，应该是(b*t)个patch，每个patch有p个字节
        bytes_input = rearrange(bytes, "b (t p) -> (b t) p", p=self.patch_size)

        # padding_local和padding_global类似，形状是 (b*t, p)
        padding_local = bytes_input.new(bytes_input.shape[0], 1).fill_(self.pad)

        # 给local的形状，应该还是(b*t,p)，扔掉最后一个，前面插入一个padding
        bytes_local = torch.cat((padding_local, bytes_input[:, :-1]), -1)

        return bytes_global, bytes_local

MegaByte 是一个可扩展的架构，它使 Transformer 大模型能够支持百万级的上下文，并且生成速度也更快。但是目前对 MegaByte 的测试规模还太小，说它能取代经典 Transformer 大模型还为时过早。

另外，序列里如果有连续的空格或者回车，这些也是按原样切成 patch，还是多个空格只算一个呢？我的猜想应该是按原样，毕竟在 Python 里空格个数是很重要的…

论文地址：https://arxiv.org/abs/2305.07185