gemma-2B-10M：10M上下文递归局部注意力

gemma-2B-10M简介

gemma-2B-10M 是一个具有高达 10M 上下文长度的递归局部注意力模型。gemma-2B-10M的实现仅使用 小于32GB 的内存！

Gemma 2B特性：

• • gemma-2B-10M 模型支持 10M（1000万） Token上下文序列长度。

• • gemma-2B-10M 模型可以在少于 32GB 的内存上可推理运行。

• • gemma-2B-10M 模型为 CUDA 做了本地推理优化。

• • 采用了递归局部注意力，实现 O(N) 内存复杂度。

gemma-2B-10M工作原理

对于LLM模型上下文扩展的最大瓶颈是 KV 缓存的大小，这涉及在计算最新令牌的注意力之前存储先前令牌的键值对。如果不这样做，计算成本会呈二次方增长，对于更长的序列这几乎是必须的。下面的 GIF 图解说明了 KV 缓存避免重新计算的概念。

递归局部注意力

gemma-2B-10M采用了最新 Transformer-XL 和递归神经网络（recurrent neural networks）的方法，改为在局部 2048 x 2048 块上计算注意力，并将该块传递给 MLP 以存储状态向量。然后，这个状态向量作为额外的参数传递到计算的下一个局部注意力块中，目标是保留从早期令牌序列中的记忆。

Infini-Attention

同时，在建立了递归状态向量的实用性之后，借鉴自 Google 的开创性论文中 Infini-attention  方法，使用压缩内存来存储先前层的信息。这允许模型只需要对局部块执行标准的点积（二次方注意力），并线性地关注过去的压缩内存。

逐步增加上下文窗口大小

最后，为了优化训练成本和数据，gemma-2B-10M遵循了 GrowLength 中的方法，逐渐将上下文大小从 32K 逐步增加，直到 10M上下文窗口（32K → 64K → 128K → 256K → 512K → 1M → 2M → 4M → 10M）。这允许研发团队在开始时优先使用较短的序列进行预训练，从而提供更高的利用率，并允许模型学习更简单的表示嵌入。随着上下文窗口的扩展，随后从更强大的初始表示状态中受益，这简化了更大上下文窗口大小的训练。

• • InfiniAttention：https://arxiv.org/pdf/2404.07143

• • Transformer-XL：https://arxiv.org/pdf/1901.02860

注意： 当前开放的Gemma 2B 模型属于早期模型权重，仅进行了 200 步的模型训练。模型下载地址：https://huggingface.co/mustafaaljadery/gemma-2B-10M

附录

• • github：https://github.com/mustafaaljadery/gemma-2B-10M?tab=readme-ov-file

• • huggingface：https://huggingface.co/mustafaaljadery/gemma-2B-10M

• • 官方博客：https://aksh-garg.medium.com/gemma-10m-technical-overview-900adc4fbeeb