simple-hierarchical-transformer - 分层Transformer模型探索多层次预测编码

简单层次化 Transformer

围绕在 GPT 中引入多个层次化预测编码模型的简单想法进行的实验。这个想法非常简单,可能不会奏效。但是,深度学习的进步正是建立在简单想法的基础之上。值得一试。

到目前为止,这个想法已经通过了一位研究朋友的试金石测试。我将在接下来的一周左右完成它。如果不成功,我也会保留负面的实验结果和代码库,也许有博士生可以在此基础上进行改进。

更新:我认为它正在奏效 🤞

致谢

感谢 StabilityAI 赞助进行这项独立研究
感谢 🤗 Huggingface 提供的 accelerate 库

安装

$ pip install simple-hierarchical-transformer

使用方法

三个层次结构,全部用于预测下一个标记

import torch
from simple_hierarchical_transformer import HierarchicalTransformer

model = HierarchicalTransformer(
    num_tokens = 20000,                # 标记数量
    dim = 512,                         # 模型维度
    depth = 6,                         # 深度
    dim_head = 64,                     # 每个注意力头的维度
    heads = 8,                         # 注意力头数量
    seq_len = 2048,                    # 序列长度
    hierarchies = (1, 2, 8),           # 层次结构 - 这里我们有 1x (像常规 Transformer), 然后是 2x 和 8x 压缩的层次标记,它们经过各自的 Transformer 块。在每一层,信息被合并到一个层次结构中
    window_sizes = (32, 64, None)      # 局部注意力窗口大小 - 想法是更高层次可以将远距离信息传递给局部层次。None 表示完整的接受域。设置为 0 将完全关闭该层次的注意力(而标记移位仍然会在每一层生效)
)

ids = torch.randint(0, 20000, (1, 2048))

loss, _ = model(ids, return_loss = True)
loss.backward()

# 经过大量训练后

logits = model(ids)

如果不指定 hierarchies 和 window_sizes,你基本上就是默认使用一个具有全序列长度注意力的常规自回归 Transformer。


# 非层次化 Transformer

model = HierarchicalTransformer(
    num_tokens = 20000,
    dim = 512,
    depth = 8,
    dim_head = 64,
    heads = 8,
    seq_len = 2048,
    hierarchies = 1,        # 如果不设置则默认为 1
    window_sizes = None     # 如果不设置则默认为 None (全序列长度)
)

现在是一个更复杂的例子。实验表明,随着层次结构的压缩,你需要更大的模型维度来获得适当的容量。

model = HierarchicalTransformer(
    num_tokens = 256,
    dim = (128, 256, 512, 1024),
    depth = 8,
    seq_len = 1024,
    use_flash_attn = True,
    ff_mult = (2, 2, 4, 4),
    dim_head = (16, 32, 64, 64),
    heads = (2, 4, 8, 8),
    hierarchies = (1, 2, 4, 16),
    hierarchical_stride = (1, 1, 1, 8),  # 这将决定压缩时的步长,以及将层次标记连接到精细标记时,过去的标记将重复这么多次。因果性不会被违反,因为使用了沙漏 Transformer 的技巧,即序列被压缩因子 - 1 移动。建议除了高度压缩的层次外,保持为 1,因为它会变得非常不具竞争力,且生成结果看起来有问题
    window_sizes = (16, 32, 64, None)
).cuda()

# 层次结构
# 1x - 维度 128 - 注意力 (2 个头, 16 维, 接受域 16)
# 2x - 维度 256 - 注意力 (4 个头, 32 维, 接受域 32)
# 4x - 维度 512 - 注意力 (8 个头, 64 维, 接受域 64)
# 8x - 维度 1024 - 注意力 (8 个头, 64 维, 全部接受域)

待办事项

引用

最接近的想法是沙漏型变换器。

我对层级方法重新产生兴趣是源于阅读这篇文章。

@article{Nawrot2021HierarchicalTA,
    title   = {Hierarchical Transformers Are More Efficient Language Models},
    author  = {Piotr Nawrot and Szymon Tworkowski and Michal Tyrolski and Lukasz Kaiser and Yuhuai Wu and Christian Szegedy and Henryk Michalewski},
    journal = {ArXiv},
    year    = {2021},
    volume  = {abs/2110.13711}
}

@inproceedings{dao2022flashattention,
    title   = {Flash{A}ttention: Fast and Memory-Efficient Exact Attention with {IO}-Awareness},
    author  = {Dao, Tri and Fu, Daniel Y. and Ermon, Stefano and Rudra, Atri and R{\'e}, Christopher},
    booktitle = {Advances in Neural Information Processing Systems},
    year    = {2022}
}

@misc{su2021roformer,
    title   = {RoFormer: Enhanced Transformer with Rotary Position Embedding},
    author  = {Jianlin Su and Yu Lu and Shengfeng Pan and Bo Wen and Yunfeng Liu},
    year    = {2021},
    eprint  = {2104.09864},
    archivePrefix = {arXiv},
    primaryClass = {cs.CL}
}

@inproceedings{Sun2022ALT,
    title     = {A Length-Extrapolatable Transformer},
    author    = {Yutao Sun and Li Dong and Barun Patra and Shuming Ma and Shaohan Huang and Alon Benhaim and Vishrav Chaudhary and Xia Song and Furu Wei},
    year      = {2022}
}

@software{peng_bo_2021_5196578,
    author    = {PENG Bo},
    title     = {BlinkDL/RWKV-LM: 0.01},
    month     = {aug},
    year      = {2021},
    publisher = {Zenodo},
    version   = {0.01},
    doi       = {10.5281/zenodo.5196578},
    url       = {https://doi.org/10.5281/zenodo.5196578}
}

@article{Piergiovanni2023Mirasol3BAM,
    title   = {Mirasol3B: A Multimodal Autoregressive model for time-aligned and contextual modalities},
    author  = {A. J. Piergiovanni and Isaac Noble and Dahun Kim and Michael S. Ryoo and Victor Gomes and Anelia Angelova},
    journal = {ArXiv},
    year    = {2023},
    volume  = {abs/2311.05698},
    url     = {https://api.semanticscholar.org/CorpusID:265129010}
}