PyTorch实现 Self Attention

点击上方“小白学视觉”，选择加"星标"或“置顶”重磅干货，第一时间送达通过修改SelfAttention的执行逻辑，可以节省大量的激活值显存开销。这篇文章的消除方法来自于2021年12月10日谷歌放到arxiv上的文章self attention does not need O(n^2) memory. 该方法巧妙地使用了小学学到的加法分配率，将self attention中的固定激活值降到了O(1)的程度。[1]Self Attention 固定激活值显存分析Hugging face Transformers中，SelfAttention 内核实现表格中只列举了会实测中产生激活值的操作，其中B为Batch_size，L为sequence_length，H为hidden_size，m为SelfAttention中head的数量。则总和 。观察：当  固定时, 即模型结构是固定的时候, 我们发现激活值是和  线性相关的。当  变化时, 我们发现会存在一个常数项 , 我称这个常数激活值开销为固定激活值。这个主要是在Query和Key矩阵做乘法, 以及后续的一 ………………………………