LongNet：将Transformer扩展至10亿个标记

这是论文《LongNet：将Transformer扩展至10亿个标记》（作者：丁家宇、马书铭、董莉、张星星、黄绍涵、王文辉、魏福鲁）的开源实现。LongNet是一种Transformer变体，旨在将序列长度扩展到超过10亿个标记，同时不牺牲短序列上的性能。

安装

pip install longnet

使用方法

安装LongNet后，您可以按如下方式使用DilatedAttention类：

import torch
from long_net import DilatedAttention


# 模型配置
dim = 512
heads = 8
dilation_rate = 2
segment_size = 64

# 输入数据
batch_size = 32
seq_len = 8192


# 创建模型和数据
model = DilatedAttention(dim, heads, dilation_rate, segment_size, qk_norm=True)
x = torch.randn((batch_size, seq_len, dim))

output = model(x)
print(output)

LongNetTransformer

一个完全可训练的Transformer模型，包含带有层归一化的前馈网络、SWIGLU以及并行Transformer块的扩张注意力层。

import torch
from long_net.model import LongNetTransformer

longnet = LongNetTransformer(
    num_tokens=20000,
    dim=512,
    depth=6,
    dim_head=64,
    heads=8,
    ff_mult=4,
)

tokens = torch.randint(0, 20000, (1, 512))
logits = longnet(tokens)
print(logits)

训练

• 要在enwiki8数据集上运行一个简单的训练任务，请先克隆仓库，安装requirements.txt中的依赖项，然后运行python3 train.py。

LongNet 简介

在大型语言模型时代，序列长度的扩展已成为关键瓶颈。然而，现有方法要么面临计算复杂度问题，要么受限于模型表达能力，从而限制了最大序列长度。本文提出了LongNet，一种能够将序列长度扩展到超过10亿个标记的Transformer变体，且不会降低短序列上的性能。具体而言，他们提出了扩张注意力机制，该机制随着距离的增加呈指数级扩大注意力范围。

特点

LongNet具有显著优势：

1. 其计算复杂度为线性，且与标记数呈对数关系。
2. 可作为分布式训练器用于极长序列。
3. 其扩张注意力机制可以直接替代标准注意力机制，能够无缝集成到现有的基于Transformer的优化框架中。

实验结果表明，LongNet在长序列建模和通用语言任务上均表现出色。他们的工作为处理超长序列开辟了新的可能性，例如将整个语料库甚至整个互联网视为一个序列来处理。

引用

@inproceedings{ding2023longnet,
  title={LongNet: Scaling Transformers to 1,000,000,000 Tokens},
  author={Ding, Jiayu and Ma, Shuming and Dong, Li and Zhang, Xingxing and Huang, Shaohan and Wang, Wenhui and Wei, Furu},
  booktitle={Proceedings of the 10th International Conference on Learning Representations},
  year={2023}
}

---

待办事项

• 修复ParallelTransformer Block的前向传播，使其使用扩张注意力
• 在enwiki 8数据集上进行训练并测试
• 实现多头注意力机制

LongNet 快速上手指南

LongNet 是一种 Transformer 变体，旨在将序列长度扩展至 10 亿（1,000,000,000）个 token，同时不牺牲短序列上的性能。其核心创新在于“膨胀注意力机制”（Dilated Attention），能够以线性计算复杂度处理超长序列。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux、macOS 或 Windows
• Python 版本：建议 Python 3.8 及以上
• 前置依赖：
  • torch (PyTorch)
  • pip 包管理工具

提示：国内用户建议使用清华源或阿里源加速 pip 安装过程。

安装步骤

使用 pip 直接安装 LongNet：

pip install longnet

国内加速安装（推荐）：

pip install longnet -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

LongNet 提供了两种主要的使用方式：直接使用核心的 DilatedAttention 模块，或使用封装好的 LongNetTransformer 模型。

1. 使用 DilatedAttention 模块

这是最基础的用法，适用于将膨胀注意力机制集成到现有架构中。

import torch
from long_net import DilatedAttention


# 模型配置
dim = 512
heads = 8
dilation_rate = 2
segment_size = 64

# 输入数据
batch_size = 32
seq_len = 8192


# 创建模型和数据
model = DilatedAttention(dim, heads, dilation_rate, segment_size, qk_norm=True)
x = torch.randn((batch_size, seq_len, dim))

output = model(x)
print(output)

2. 使用 LongNetTransformer 完整模型

这是一个包含前馈网络、层归一化、SWIGLU 激活函数以及并行 Transformer 块的即训即用模型。

import torch
from long_net.model import LongNetTransformer

longnet = LongNetTransformer(
    num_tokens=20000,
    dim=512,
    depth=6,
    dim_head=64,
    heads=8,
    ff_mult=4,
)

tokens = torch.randint(0, 20000, (1, 512))
logits = longnet(tokens)
print(logits)

3. 简单训练示例

若需在 enwiki8 数据集上进行简单的训练测试，请克隆仓库并安装依赖后运行：

git clone <repository_url>
cd LongNet
pip install -r requirements.txt
python3 train.py