尾部火炬

Zhai等人的Attention Free Transformer层的非官方 PyTorch 实现。[ abs , pdf ] 来自 Apple Inc.

安装

您可以aft_pt通过以下方式安装pip：

pip install aft_pt

用法

您可以从包中导入AFT-Full或AFT-Simple层（如本文所述），如下所示：

AFTFull

from aft_pt import AFTFull

layer = AFTFull(
    max_seqlen=20,
    dim=512,
    hidden_dim=64
)

# a batch of sequences with 10 timesteps of length 512 each
x = torch.rand(32, 10, 512)
y = layer(x) # [32, 10, 512]

AFTSimple

from aft_pt import AFTSimple

layer = AFTSimple(
    max_seqlen=20,
    dim=512,
    hidden_dim=64
)

# a batch of sequences with 10 timesteps of length 512 each
x = torch.rand(32, 10, 512)
y = layer(x) # [32, 10, 512]

AFTLocal

from aft_pt import AFTLocal

layer = AFTLocal(
    max_seqlen=20,
    dim=512,
    hidden_dim=64
)

# a batch of sequences with 10 timesteps of length 512 each
x = torch.rand(32, 10, 512)
y = layer(x) # [32, 10, 512]

该层包装器是您现有网络/变形金刚的“即插即用”。您可以将 Self-Attention 层换成此包中的可用层，只需进行最少的更改。

去做

• 添加完整的 AFT 架构
• 添加变体，例如，AFTConv
• 使用 Karpathy 的minGPT进行基准测试

贡献

如果你喜欢这个repo，请给我一个star！如果有任何修改或建议，请随时提出 PR/问题。

学分

@misc{attention-free-transformer,
title = {An Attention Free Transformer},
author = {Shuangfei Zhai and Walter Talbott and Nitish Srivastava and Chen Huang and Hanlin Goh and Ruixiang Zhang and Josh Susskind},
year = {2021},
URL = {https://arxiv.org/pdf/2105.14103.pdf}
}

执照

麻省理工学院