MoCo中Key Encoder反向传播难处理的真正原因是什么？

MoCo原论文提到：使用队列可扩大字典规模，但会导致Key Encoder的反向传播难以处理（梯度需传播至队列中所有样本）。
起初我认为核心原因是队列操作不可微分，但实际可计算队列内所有样本的梯度，反向传播可正常执行（见下方代码）。那么Key Encoder反向传播难处理的真正原因是什么？我猜测可能是队列（字典）规模过大导致内存激增。

q = nn.Linear(768,128)
k = nn.Linear(768,128)
bs = 64
ks = 4095
model = nn.ModuleList([q,k])
x = torch.randn(bs, 768)
optim = torch.optim.SGD(model.parameters(),lr=0.01)
loss = nn.CrossEntropyLoss()
def forward(x):
    xq = q(x)
    xk = k(x + 0.1)
    que = torch.rand(ks,128)
    pos = torch.einsum("nc,nc->n",xq,xk)
    neg = torch.einsum("nc,kc->nk",xq,que)
    out = torch.cat([pos.unsqueeze(-1),neg],dim=1)
    t = torch.zeros(out.shape[0],dtype=torch.long)
    l = loss(out,t)
    return l
loss = forward(x)
loss.backward()
optim.step()

你的猜测完全正确，MoCo里Key Encoder反向传播难处理的核心问题就是队列规模过大导致内存爆炸。

你提供的代码里队列是随机生成的独立张量，和Key Encoder没有计算图关联，所以反向传播时不会对队列内容求导。但在真实的MoCo训练流程中，队列里的每个key都是Key Encoder对历史输入样本的输出——也就是说，队列内所有key都和Key Encoder的参数存在计算图依赖关系。

当队列规模达到比如65536这样的量级时，反向传播需要同时维护这数万样本对应的计算图，还要计算每个样本对Key Encoder参数的梯度，这会瞬间占用极大显存，普通GPU根本无法承载。

MoCo的解决方案是不对队列中的历史key做反向传播：只让当前批次的key参与梯度计算，队列里的旧key直接用动量更新后的Key Encoder输出替换，相当于切断了旧key和当前Key Encoder参数的计算图关联，以此大幅节省内存。

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