摘要
我们提出Spartan，一种用于训练具有预定稀疏度水平的稀疏神经网络模型的方法。Spartan基于两种技术的结合：（1）通过正则化最优传输问题实现对低幅值参数的软Top-k掩码；（2）在前向传播中采用基于双平均（dual averaging）的参数更新机制，并引入硬性稀疏化。该框架实现了探索与利用之间的权衡：在训练初期，学习器能够探索多种稀疏模式；随着训练的进行，软Top-k近似逐渐趋于尖锐化，系统逐渐从探索转向在固定稀疏掩码下的参数优化。Spartan具有足够的灵活性，可支持多种稀疏分配策略，包括非结构化稀疏、块状结构稀疏，以及由参数级成本线性模型所调控的一般成本敏感型稀疏分配。在ImageNet-1K分类任务中，Spartan可生成稀疏度达95%的ResNet-50模型和稀疏度达90%的块状稀疏ViT-B/16模型，与全连接训练相比，其Top-1准确率绝对损失均低于1%。