对PosWiseFFN的改进： MoE、PKM、UltraMem

先从PosWiseFFN说起

class PoswiseFeedForwardNet(nn.Module):
    def __init__(self):
        super(PoswiseFeedForwardNet, self).__init__()
        self.fc = nn.Sequential(
            nn.Linear(d_model, d_ff, bias=False),
            nn.GeLU(),
            nn.Linear(d_ff, d_model, bias=False))

    def forward(self, inputs):                                  # inputs: [batch_size, seq_len, d_model]
        residual = inputs
        output = self.fc(inputs)
        return nn.LayerNorm(d_model)(output + residual)  # [batch_size, seq_len, d_model]

如果Attention的维度是d_model，通常PosWiseFFN模型结构就是2个矩阵中间加个Gelu，d_ff是d_model的4倍：第1个矩阵的weight是[d_model, 4*d_model]，第2个矩阵的的weight是[4*d_model, d_model]。

PosWiseFFN这个结构也可以理解成一种qkv查询的思路，如果第1个矩阵理解成key，第二矩阵理解成value，那么输入就是[batch_size, seq_len, d_model]的input作为query先去和key做矩阵乘法，得到一个[batch_size, seq_len, 4*d_model]的dots，这个dots过了GeLU后再去和[4*d_model, d_model]的第二个矩阵相乘，这一步变向取了前d_model重要的结果。问题来了，能不能把 4*d_model的d_ff给变得更大呢？Figure 1来自Large Memory Layers with Product Keys的Figure1，图里的|K|在PosWiseFFN里就是 4*d_model。

下面的PKM简单来说就是把这种qkv查询的思路借用PQ的思想给改进了

PKM（Product Key Memory，这个Product其实就是Product Quantization的Product）

在Large Memory Layers with Product Keys的Figure1里，q的shape是[…,d_model]，k的shape是[d_model, |K|]，下面看Figure2里怎么解决|K|过大的问题？图里把d_model维的q劈成q1和q2，q1和q2的维度分别是d_model/2；同样的，把[d_model, |K|]的keys劈成[d_model/2, |K|]的sub-key set 1（下图里不带’的$c_1$,$c_2$,$c_3$）和[d_model/2, |K|]的sub-key set 2（下图里带’的$c_1^{}$