如何在不复制数据的情况下将MultiheadAttention多头应用于同一输入？

你的核心需求是让多个注意力头基于同一序列计算注意力，但避免复制原始数据带来的内存浪费。原来通过repeat复制数据的方式确实不够高效，这里提供两种更优的实现方式：

方案一：手动实现多头注意力逻辑 #

直接基于原始序列构建多注意力头的查询、键、值投影，全程无需复制数据：

import torch
import torch.nn as nn

N, C, T = 2, 3, 5
n_heads = 7
X = torch.rand(N, T, C)

# 定义投影层：将原始C维特征映射到n_heads个C维特征（总维度C*n_heads）
q_proj = nn.Linear(C, C * n_heads)
k_proj = nn.Linear(C, C * n_heads)
v_proj = nn.Linear(C, C * n_heads)
out_proj = nn.Linear(C * n_heads, C * n_heads)  # 可选，和原生MultiheadAttention对齐

# 生成Q/K/V并拆分注意力头
Q = q_proj(X).view(N, T, n_heads, C).transpose(1, 2)  # shape: (N, n_heads, T, C)
K = k_proj(X).view(N, T, n_heads, C).transpose(1, 2)
V = v_proj(X).view(N, T, n_heads, C).transpose(1, 2)

# 计算注意力分数
attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / (C ** 0.5)
attn_probs = torch.softmax(attn_scores, dim=-1)

# 计算输出并合并注意力头
output = torch.matmul(attn_probs, V).transpose(1, 2).flatten(2)  # shape: (N, T, C*n_heads)
output = out_proj(output)  # 可选，应用输出投影

这个方案通过一次线性投影生成所有注意力头的Q/K/V，再通过维度拆分和重组实现多头计算，完全没有复制原始序列数据。

方案二：复用原生MultiheadAttention但避免数据复制 #

如果你想继续使用torch.nn.MultiheadAttention，可以通过调整输入的投影方式，替代数据复制：

import torch
import torch.nn as nn

N, C, T = 2, 3, 5
n_heads = 7
X = torch.rand(N, T, C)

# 先将原始特征投影到C*n_维，无需复制数据
proj = nn.Linear(C, C * n_heads)
X_proj = proj(X)  # shape: (N, T, C*n_heads)

# 使用原生MultiheadAttention，此时embed_dim=C*n_heads，num_heads=n_heads
attn = nn.MultiheadAttention(C * n_heads, n_heads, batch_first=True)
output, _ = attn(X_proj, X_proj, X_proj)

这种方式和你原来的效果一致，但通过线性投影替代了数据复制，内存占用更低（原始X只存储一次，而非n_heads次）。

两种方案的核心思路都是：用线性投影生成多注意力头所需的特征，而非复制原始数据，这样既满足了多头计算的需求，又避免了不必要的内存开销。

内容的提问来源于stack exchange，提问作者EZLearner