partial self-attention

### 部分自注意力机制概述

部分自注意力机制是一种改进版的自注意力模型，旨在减少计算复杂度并提高处理效率。传统自注意力机制会考虑序列中的每一个位置与其他所有位置之间的关系，而部分自注意力则只关注局部区域内的交互作用[^1]。

这种设计使得网络能够在保持性能的同时显著降低资源消耗，在大规模数据集上展现出更好的扩展性和实用性。具体来说，通过限制查询向量与键向量之间的作用范围，可以有效控制内存占用和运算时间。

### 实现方式

为了实现部分自注意力机制，通常采用滑动窗口方法来限定每个token所能访问到的历史上下文长度。以下是基于PyTorch框架的一个简单示例：

import torch
from torch import nn

class PartialSelfAttention(nn.Module):
    def __init__(self, embed_size, heads=8, window_size=3):
        super(PartialSelfAttention, self).__init__()
        
        assert (embed_size % heads == 0), "Embedding size must be divisible by number of heads"
        
        self.embed_size = embed_size
        self.heads = heads
        self.window_size = window_size
        
        self.values = nn.Linear(self.embed_size, self.embed_size, bias=False)
        self.keys = nn.Linear(self.embed_size, self.embed_size, bias=False)
        self.queries = nn.Linear(self.embed_size, self.embed_size, bias=False)
        self.fc_out = nn.Linear(self.embed_size, self.embed_size)

    def forward(self, values, keys, query, mask=None):
        N = query.shape[0]
        value_len, key_len, query_len = values.shape[1], keys.shape[1], query.shape[1]

        # Split embedding into self.heads different pieces
        values = values.reshape(N, value_len, self.heads, self.embed_size // self.heads)
        keys = keys.reshape(N, key_len, self.heads, self.embed_size // self.heads)
        queries = query.reshape(N, query_len, self.heads, self.embed_size // self.heads)

        energy = torch.einsum("nqhd,nkhd->nhqk", [queries, keys])

        if mask is not None:
            energy = energy.masked_fill(mask == 0, float("-1e20"))

        attention = torch.softmax(energy / (self.embed_size ** (1 / 2)), dim=3)

        out = torch.einsum("nhql,nlhd->nqhd", [attention, values]).reshape(
            N, query_len, self.embed_size
        )

        out = self.fc_out(out)
        return out

在这个例子中，`window_size`参数用于定义每个query能够看到的最大距离。当设置较小值时，即实现了所谓的“局部化”的自我注意功能；如果将其设为整个输入序列长度，则退化成标准形式[^2]。

### 应用场景

部分自注意力机制广泛应用于自然语言处理领域之外的任务当中，特别是在通信系统优化方面表现出色。例如，在光通信干扰消除研究里提到利用深度学习技术解决非正交信号带来的多址接入问题时，就采用了类似的思路来构建更高效的编码解码器结构。