[ACM MM 2024] AdapMTL: 多任务学习模型的自适应剪枝框架

AdapMTL: Adaptive Pruning Framework for Multitask Learning Model

AdapMTL: Adaptive Pruning Framework for Multitask Learning Model (arxiv.org)

Abstract

In the domain of multimedia and multimodal processing, the efficient handling of diverse data streams such as images, video, and sensor data is paramount. Model compression and multitask learning (MTL) are crucial in this field, offering the potential to address the resource-intensive demands of processing and interpreting multiple forms of media simultaneously. However, effectively compressing a multitask model presents significant challenges due to the complexities of balancing sparsity allocation and accuracy performance across multiple tasks. To tackle the challenges, we propose AdapMTL, an adaptive pruning framework for MTL models. AdapMTL leverages multiple learnable soft thresholds independently assigned to the shared backbone and the task-specific heads to capture the nuances in different components’ sensitivity to pruning. During training, it co-optimizes the soft thresholds and MTL model weights to automatically determine the suitable sparsity level at each component to achieve both high task accuracy and high overall sparsity. It further incorporates an adaptive weighting mechanism that dynamically adjusts the importance of task-specific losses based on each task’s robustness to pruning. We demonstrate the effectiveness of AdapMTL through comprehensive experiments on popular multitask datasets, namely NYU-v2 and Tiny-Taskonomy, with different architectures, showcasing superior performance compared to state-of-the-art pruning methods.

背景：在多媒体和多模态处理领域，高效处理各种数据流（如图像、视频和传感器数据）具有重要意义。模型压缩和多任务学习（MTL）是该领域的关键技术，具备同时处理和解释不同媒体形式资源密集型需求的潜力。

问题：然而，有效压缩多任务模型面临着重大挑战，因为需要平衡稀疏度分配与准确性表现之间的权衡。

方法：为了应对这些挑战，本文提出了AdapMTL框架，一种自适应剪枝方法专门针对MTL模型设计而来。AdapMTL利用可学习的软阈值来分配给共享骨干和特定任务头，并捕捉到不同组件对剪枝敏感性微小差异。在训练过程中，它通过优化软阈值和MTL模型权重来确定每个组件合适的稀疏度水平，从而实现高任务准确性和整体稀疏度。此外，在特定任务损失上采用自适应加权机制以动态调整其对剪枝的鲁棒性重要性。

结果：本文在流行的多任务数据集（如NYU-v2和Tiny-Taskonomy）上进行了全面实验，并展示了AdapMTL相较于最先进剪枝方法更卓越的性能表现。

研究动机

在多媒体和多模态处理领域，高效处理图像、视频和传感器数据等多样化的数据流至关重要。模型压缩和多任务学习（MTL）是这一领域的关键技术，它们能够解决同时处理和解释多种媒体形式的资源密集型需求。然而，有效压缩多任务模型面临重大挑战，尤其是如何在多个任务之间平衡稀疏度分配和准确性表现。传统的剪枝技术主要针对单任务模型，而多任务模型具有任务间的相互依赖性、表示共享以及任务头对剪枝的不同敏感度等复杂性，使得直接应用现有剪枝方法变得困难。

AdapMTL的研究动机在于：

共享骨干和任务特定头对剪枝的敏感度不同：当前的方法没有充分考虑到这一点，导致在剪枝过程中对所有组件一视同仁，未能根据它们的不同敏感度进行优化。

训练损失的变化可以作为分配稀疏性的有用指南：对于训练损失稳定的任务，可以更激进地进行剪枝，因为它们对剪枝的鲁棒性更强；相反，如果某个任务的损失波动显著，则应减少剪枝，以避免训练难以在高稀疏度水平下收敛。

为了应对这些挑战，研究者提出了AdapMTL，一个自适应剪枝框架，用于多任务学习模型。

方法

AdapMTL的方法主要包括以下几个步骤：

问题定义：
将多任务模型剪枝定义为一个优化问题，目标是在给定数据集和期望稀疏度水平下，找到一组稀疏权重，以最小化所有任务的损失之和。数学上，这可以表述为式（1）的优化问题。

自适应剪枝算法：

 + 学习软阈值：AdapMTL使用多个可学习的软阈值，分别分配给共享骨干和每个任务特定的头。这些软阈值在训练过程中与模型权重共同优化，以自动确定每个组件的合适稀疏度水平。

+ 前向传播：在前向传播中，只有大于相应阈值的权重才会计入模型，其他权重设置为零。

+ 反向传播：在反向传播中，自动更新所有组件的阈值，以平滑地引入稀疏性。

自适应加权机制：

根据每个任务对剪枝的鲁棒性，动态调整任务特定损失的重要性。如果某个任务的训练损失稳定，说明它对剪枝较为鲁棒，可以赋予更高的权重；反之，如果损失波动大，则应减少其权重。

训练过程：

AdapMTL不需要任何预训练或预剪枝的模型，可以从头开始训练。在训练过程中，模型权重和组件级软阈值同时更新，以实现高稀疏度和高任务准确性的平衡。