sinkhorn-transformer:Sinkhorn变压器-稀疏Sinkhorn注意的实际实现

灌角变压器 这是对“概述的工作的复制，并进行了其他增强。 它包括一个参数化的排序网络，使用接收器喇叭归一化对与最相关的关键字存储桶与查询存储桶相匹配的排列矩阵进行采样。 这项工作还引入了可逆网络和前馈分块（ 引入的概念），以进一步节省内存。 204k令牌（演示目的） 安装 $ pip install sinkhorn_transformer 用 基于Sinkhorn变压器的语言模型 import torch from sinkhorn_transformer import SinkhornTransformerLM model = SinkhornTransformerLM ( num_tokens = 20000 , dim = 1024 , heads = 8 , depth = 12 , max_seq_len = 8192 ,