factorised self-attention

### 因子化自注意力机制的概念与实现 #### 概念解释 因子化自注意力（Factorized Self-Attention, FSA）是一种优化方法，旨在减少标准多头自注意力层中的计算复杂度和内存消耗。传统上，在Transformer模型中使用的multi-head self-attention通过多个独立的self-attention操作来捕捉输入序列的不同方面特征[^1]。 然而，随着序列长度的增长，这种做法会导致显著增加的时间成本以及更高的资源需求。为了缓解这些问题，FSA引入了一种新的架构设计思路——将原始的大规模矩阵运算分解成更小的部分处理单元。具体而言： - **键空间降维**：通过对键向量应用线性变换以降低其维度； - **值空间压缩**：同样地对值向量执行相似的操作； - **共享参数**：允许不同头部之间部分或全部共享权重参数； 这些措施共同作用下可以有效减小整体网络尺寸而不明显牺牲性能表现。 #### 实现方式 下面给出一段基于PyTorch框架下的简化版因子化自注意力建模代码片段作为参考： ```python import torch.nn as nn import torch class FactorizedSelfAttention(nn.Module): def __init__(self, d_model, n_heads=8, reduction_ratio=4): super(FactorizedSelfAttention, self).__init__() # 定义投影映射函数 self.query_proj = nn.Linear(d_model, d_model) self.key_proj = nn.Linear(d_model, int(d_model / reduction_ratio)) self.value_proj = nn.Linear(d_model, int(d_model / reduction_ratio)) self.n_heads = n_heads def forward(self, queries, keys, values): Q = self.query_proj(queries) # [batch_size, seq_len_q, d_model] K = self.key_proj(keys) # [batch_size, seq_len_k, reduced_dim] V = self.value_proj(values) # [batch_size, seq_len_v, reduced_dim] batch_size, seq_len_q, _ = Q.size() _, seq_len_k, _ = K.size() # 将Q,K,V按照head切分并转置 Q_split = Q.view(batch_size, seq_len_q, self.n_heads, -1).transpose(1, 2) K_split = K.view(batch_size, seq_len_k, self.n_heads, -1).transpose(1, 2) V_split = V.view(batch_size, seq_len_k, self.n_heads, -1).transpose(1, 2) scores = torch.matmul(Q_split, K_split.transpose(-2,-1)) / np.sqrt(K_split.shape[-1]) attention_weights = torch.softmax(scores,dim=-1) output = torch.matmul(attention_weights, V_split) # 合并heads并将结果投回原维度 combined_output = output.transpose(1, 2).contiguous().view(batch_size, seq_len_q, -1) return combined_output ``` 上述代码实现了基本的功能模块，实际应用场景可能还需要考虑更多细节调整如残差连接、层归一化等组件加入其中。