特征交叉系列：DCN-Mix 混合低秩交叉网络理论和实践

关键词：DCN，DCN-Mix，推荐算法，MOE混合专家网络，低秩矩阵分解

内容摘要

• DCN-Mix和DCN-V2的关系
• DCN-V2权重矩阵的低秩性和矩阵分解
• DCN-Mix的混合专家网络
• DCN-Mix在PyTorch下的实践
• DCN-Mix调参和效果对比

DCN-Mix和DCN-V2的关系

DCN-Mix(a mixture of low-rank DCN)是基于DCN-V2的改进版，它提出使用矩阵分解来降低DCN-V2的时间空间复杂度，又引入多次矩阵分解来达到类似混合专家网络MOE的效果从而提升交叉层的表征能力，若读者对DCN-V2不甚了解可以参考上一节特征交叉系列：Deep&Cross(DCN-V2)理论和实践做知识铺垫。

---

DCN-V2权重矩阵的低秩性和矩阵分解

在DCN-V2中核心的参数是交叉层的权重矩阵W，该参数是M×M的方阵，其中M是所有输入embedding拼接后的向量总长度，每一层交叉之间W不共享，W矩阵需要学习的参数数量能占到所有参数量的70%以上，而进一步作者发现随着网络的训练，W矩阵的奇异值出现快速下降呈现出低秩特性，代表该矩阵存在信息冗余，因此可以考虑通过矩阵分解来进行特征提取和信息压缩。
在PyTorch中可以通过torch.linalg.svd计算出矩阵的奇异值，例如

>>> a = torch.tensor([[1, 1], [1, 1.1]])
>>> u, s, v = torch.linalg.svd(a)
>>> print(s)
tensor([2.0512, 0.0488])

其中s是对角阵，斜对角线上的值就是奇异值，a矩阵的第二行几乎可以从第一行线性变换而来，因此s各位置上的奇异值差距极大，第一个奇异值基本携带了全部的矩阵信息。
在DCN-V2的训练代码里面，打印出第一个交叉层初始化的W矩阵和训练早停后W矩阵的奇异值，奇异值的长度和输入长度M一致，代码如下

# 初始化时
model = DCN(field_num=10, feat_dim=72, emb_num=16, order_num=2, dropout=0.1, method='parallel').to(DEVICE)
init_s = torch.linalg.svd(model.cross_net.cell_list[0].w)[1].cpu().detach().numpy().tolist()

# 早停时
if early_stop_flag:
    train_s = torch.linalg.svd(model.cross_net.cell_list[0].w)[1].cpu().detach().numpy().tolist()
    break

奇异值列表中元素大小逐个递减，对init_s和train_s分别做最大最小归一化，要求第一个奇异值归因化为1，

init_s = [(x - min(init_s)) / (max(init_s) - min(init_s)) for x in init_s]
train_s = [(x - min(train_s)) / (max(train_s) - min(train_s)) for x in train_s]

然后做图看一下初始矩阵的奇异值和收敛后的奇异值的各个位置元素的大小情况

import matplotlib.pylab as plt
plt.scatter(list(range(len(init_s))), init_s, label='init', s=3)
plt.scatter(list(range(len(train_s))), train_s, label='learned', s=3)
plt.legend(loc=0)
plt.show()

init和learned奇异值下降对比

相比于初始化阶段(蓝线)，模型收敛后(橙线)的W矩阵奇异值急速下降，说明头部的奇异值已经携带了大部分矩阵信息，W矩阵可以考虑做压缩。
在论文中作者将W分解为U，V两个矩阵的相乘，其中U，V都是维度为[M, R]的二维矩阵，M和输入等长，R<=M/2，公式如下

矩阵分解

此时一个交叉权重的参数数量由M平方降低为2×MR。

---

DCN-Mix的混合专家网络

DCN-Mix使用矩阵UV分解来逼近原始的交叉矩阵W，受到MOE(Mixture of Experts)混合专家网络的启发，作者对W进行多次矩阵分解，单个矩阵分解相当于单个专家网络(Expert)在子空间学习特征交叉，再引入门控机制(Gate)对多个子空间的交叉结果进行自适应地融合，从而提高交叉层的表达能力，DCN结合MOE的示意图如下

MOE示意图

其中该层的输入Input x分别进入n个Expert专家网络，专家网络中包含UV矩阵相乘，同时Input x输入给一个门控网络Gate+Softmax输出n个权重标量，最后Input x会和加权求和的专家网络结果做残差连接。
将矩阵分解和MOE结合起来形成最终的交叉层公式如下

结合MOE的矩阵分解交叉层

相比于DCN-V2，等号左侧的哈达玛积部分改为了一个Σ加权求和的UV矩阵逼近，而右侧的残差连接放到最后和MOE的结果一起做残差连接。

---

DCN-Mix在PyTorch下的实践

本次实践的数据集和上一篇特征交叉系列：完全理解FM因子分解机原理和代码实战一致，采用用户的购买记录流水作为训练数据，用户侧特征是年龄，性别，会员年限等离散特征，商品侧特征采用商品的二级类目，产地，品牌三个离散特征，随机构造负样本，一共有10个特征域，全部是离散特征，对于枚举值过多的特征采用hash分箱，得到一共72个特征。
DCN-Mix的PyTorch代码实现如下

class Embedding(nn.Module):
    def __init__(self, feat_num, emb_num):
        super(Embedding, self).__init__()
        self.embedding = nn.Embedding(feat_num, emb_num)
        nn.init.xavier_normal_(self.embedding.weight.data)

    def forward(self, x):
        # [None, filed_num] => [None, filed_num, emb_num] => [None, filed_num * emb_num]
        return self.embedding(x).flatten(1)


class DNN(nn.Module):
    def __init__(self, input_num, hidden_nums, dropout=0.1):
        super(DNN, self).__init__()
        layers = []
        input_num = input_num
        for hidden_num in hidden_nums:
            layers.append(nn.Linear(input_num, hidden_num))
            layers.append(nn.BatchNorm1d(hidden_num))
            layers.append(nn.ReLU())
            layers.append(nn.Dropout(p=dropout))
            input_num = hidden_num
        self.mlp = nn.Sequential(*layers)
        for layer in self.mlp:
            if isinstance(layer, nn.Linear):
                nn.init.xavier_normal_(layer.weight.data)

    def forward(self, x):
        return self.mlp(x)


class CrossCell(nn.Module):
    """一个交叉单元"""

    def __init__(self, input_num, r):
        super(CrossCell, self).__init__()
        self.v = nn.Parameter(torch.randn(input_num, r))
        self.u = nn.Parameter(torch.randn(input_num, r))
        self.b = nn.Parameter(torch.randn(input_num, 1))
        nn.init.xavier_normal_(self.v.data)
        nn.init.xavier_normal_(self.u.data)

    def forward(self, x0, xi):
        # [None, emb_num] => [None, emb_num, 1]
        xi = xi.unsqueeze(2)
        x0 = x0.unsqueeze(2)
        # [r, input_num] * [None, emb_num, 1] => [None, r, 1]
        # [input_num, r] * [None, r, 1] => [None, emb_num, 1]
        xii = (torch.matmul(self.u, torch.matmul(self.v.t(), xi)) + self.b) * x0
        return xii  # [None, emb_num, 1]


class MOECrossCell(nn.Module):
    def __init__(self, input_num, r, k):
        super(MOECrossCell, self).__init__()
        self.k = k
        self.cross_cell = nn.ModuleList([CrossCell(input_num, r) for i in range(self.k)])
        self.gate = nn.Linear(input_num, self.k)
        nn.init.xavier_normal_(self.gate.weight.data)

    def forward(self, x0, xi):
        # [None, emb_num] => [None, emb_num, 1]
        xii = xi.unsqueeze(2)
        export_out = []
        for i in range(self.k):
            cross_out = self.cross_cell[i](x0, xi)
            # [[None, emb_num, 1], [None, emb_num, 1], [None, emb_num, 1], [None, emb_num, 1]]
            export_out.append(cross_out)
        export_out = torch.concat(export_out, dim=2)  # [None, emb_num, 4]
        # [None, k] => [None, 1, k]
        gate_out = self.gate(xi).softmax(dim=1).unsqueeze(dim=1)
        # [None, emb_num, 4] * [None, 1, k] = [None, emb_num, k] => [None, emb_num, 1]
        out = torch.sum(export_out * gate_out, dim=2, keepdim=True)
        out = out + xii  # [None, emb_num, 1]
        return out.squeeze(2)


class CrossNet(nn.Module):
    def __init__(self, order_num, input_num, r, k):
        super(CrossNet, self).__init__()
        self.order = order_num
        self.cell_list = nn.ModuleList([MOECrossCell(input_num, r, k) for i in range(order_num)])

    def forward(self, x0):
        xi = x0
        for i in range(self.order):
            xi = self.cell_list[i](x0=x0, xi=xi)
        return xi


class DCN(nn.Module):
    def __init__(self, field_num, feat_dim, emb_num, order_num, r=16, k=4, dropout=0.1, method='parallel',
                 hidden_nums=(128, 64, 32)):
        super(DCN, self).__init__()
        input_num = field_num * emb_num
        self.embedding = Embedding(feat_num=feat_dim, emb_num=emb_num)
        self.dnn = DNN(input_num=input_num, hidden_nums=hidden_nums, dropout=dropout)
        self.cross_net = CrossNet(order_num=order_num, input_num=input_num, r=r, k=k)
        if method not in ('parallel', 'stacked'):
            raise ValueError('unknown combine type: ' + method)
        self.method = method
        linear_dim = hidden_nums[-1]
        if self.method == 'parallel':
            linear_dim = linear_dim + input_num
        self.linear = nn.Linear(linear_dim, 1)
        nn.init.xavier_normal_(self.linear.weight.data)

    def forward(self, x):
        emb = self.embedding(x)  # [None, field * emb_num]
        cross_out = self.cross_net(emb)  # [None, input_num]
        if self.method == 'parallel':
            dnn_out = self.dnn(emb)  # [None, input_num]
            out = torch.concat([cross_out, dnn_out], dim=1)
        else:
            out = self.dnn(cross_out)  # [None, input_num]
        out = self.linear(out)
        return torch.sigmoid(out).squeeze(dim=1)

在CrossCell模块中完成了一个给予UV逼近的交叉操作，在MOECrossCell模块中完成了MOE和残差连接，其中export_out和gate_out分别为专家网络的输出和门控机制的权重。
本例全部是离散分箱变量，所有有值的特征都是1，因此只要输入有值位置的索引即可，一条输入例如

>>> train_data[0]
Out[120]: (tensor([ 2, 10, 14, 18, 34, 39, 47, 51, 58, 64]), tensor(0))

x的长度为10代表10个特征域，每个域的值是特征的全局位置索引，从0到71，一共72个特征。

---

DCN-Mix调参和效果对比

对阶数(order_num)和融合策略(method)这两个参数进行调参，分别尝试1～4层交叉层，stacked和parallel两种策略，采用10次验证集AUC不上升作为早停条件，验证集的平均AUC如下

DCN调参AUC	并行parallel	串行stacked
1层交叉（2阶）	0.6345	0.6321
2层交叉（3阶）	0.6328	0.6323
3层交叉（4阶）	0.6331	0.6333
4层交叉（5阶）	0.6340	0.6331

结论依旧是parallel效果好于stacked，其中一层交叉的并行parallel达到验证集最优AUC为0.6345。
再对比一下之前文章中实践的FM，FFM，PNN，DCN-V2等一系列算法，验证集AUC和参数规模如下

算法	AUC	参数量
FM	0.6274	361
FFM	0.6317	2953
PNN*	0.6342	29953
DeepFM	0.6322	12746
NFM	0.6329	10186
DCN-parallel-3	0.6348	110017
DCN-stacked-3	0.6344	109857
DCN-Mix-parallel-1	0.6345	54501
DCN-Mix-stacked-3	0.6333	97869

使用矩阵分解逼近策略的DCN-Mix略低于原生的DCN-V2,但是还是超越一众FM系列的算法，其中以同样是三层交叉的stacked DCN为例，DCN-Mix的参数量相比于DCN-V2有所降低，也印证了论文中提到的“在模型效果和部署延迟之间找到一个平衡”。