[論文レビュー] Knowledge Distillation for Multi-task Learning

Weihong Li, Hakan Bilen|arXiv (Cornell University)|Jul 14, 2020

Domain Adaptation and Few-Shot Learning参考文献 31被引用数 5

ひとこと要約

本稿では、学習可能アダプタを介してマルチタスクネットワークの特徴をタスク固有のモデルと一致させることで、トレーニングのバランスとパフォーマンスを向上させる、知識蒸留に基づくマルチタスク学習手法を提案する。タスク固有の損失と蒸留損失を同時に最小化することで、NYU-V2 や SVHN & Omniglot などのベンチマークで、セマンティックセグメンテーション、深度推定、表面法線予測といった多様なタスクにおいて、優れたバランスの取れたパフォーマンスを達成する。

ABSTRACT

Multi-task learning (MTL) is to learn one single model that performs multiple tasks for achieving good performance on all tasks and lower cost on computation. Learning such a model requires to jointly optimize losses of a set of tasks with different difficulty levels, magnitudes, and characteristics (e.g. cross-entropy, Euclidean loss), leading to the imbalance problem in multi-task learning. To address the imbalance problem, we propose a knowledge distillation based method in this work. We first learn a task-specific model for each task. We then learn the multi-task model for minimizing task-specific loss and for producing the same feature with task-specific models. As the task-specific network encodes different features, we introduce small task-specific adaptors to project multi-task features to the task-specific features. In this way, the adaptors align the task-specific feature and the multi-task feature, which enables a balanced parameter sharing across tasks. Extensive experimental results demonstrate that our method can optimize a multi-task learning model in a more balanced way and achieve better overall performance.

研究の動機と目的

• 異なる損失の大きさと特性を持つタスクがトレーニングを支配してしまうマルチタスク学習におけるアンバランス問題に対処する。
• 高性能なタスクが低性能なタスクを圧倒しないように、すべてのタスクの全体的なパフォーマンスを向上させる。
• タスク固有のアダプタを介した特徴一致により、意味的レベルや複雑さが異なるタスク間でのパラメータ共有をバランスよく実現する。
• 事前学習済みの単一タスクモデルをマルチタスクネットワークの知識源として活用することで、より良い一般化性能と効率性を達成する。
• 損失重み付けや勾配調整手法に比べ、タスクの支配を防げないことが多くあるが、本手法はより安定的かつ制御可能な最適化プロセスを提供する。

提案手法

• オフライン段階で各タスクごとに別々のタスク固有モデルを学習し、そのパラメータを凍結する。
• マルチタスクネットワークを、タスク固有の損失と、その特徴を凍結したタスク固有モデルの特徴と一致させる蒸留損失の両方を最小化するように最適化する。
• マルチタスクネットワークの特徴を各タスク固有モデルの特徴空間に一致させるために、小さな学習可能なタスク固有アダプタを導入する。
• 線形または非線形アダプタ（例：ReLUを伴う1×1畳み込み）を用いることで、異なるタスク間での柔軟で効果的な特徴一致を実現する。
• 複数の層（例：中間層と最終共有層）に蒸留損失を適用することで、特徴一致を強化し、最適化の安定性を向上させる。
• 選択された層におけるタスク固有損失と蒸留損失の重み付き和として、エンドツーエンドのマルチタスクネットワーク学習を実行する。

実験結果

リサーチクエスチョン

• RQ1 タスク固有モデルからの知識蒸留が、マルチタスク学習におけるバランスと全体的なパフォーマンスを向上させるか？
• RQ2 単一タスクモデルの特徴と一致させることで、マルチタスクネットワークの特徴を調整することで、一般化性能が向上し、タスクの支配が軽減されるか？
• RQ3 異なるアダプタアーキテクチャ（線形対非線形）が特徴一致と最終的なパフォーマンスに与える影響は何か？
• RQ4 最後の層に限定せず、複数のネットワーク層に蒸留損失を適用することで、どのような影響があるか？
• RQ5 本手法は、既存の損失重み付けや勾配調整戦略を上回るマルチタスク最適化のバランスを達成できるか？

主な発見

• 提案手法は NYU-V2 で最先端のパフォーマンスを達成し、セマンティックセグメンテーションで 20.75 mIoU を記録し、MTAN や他のベースラインを上回った。
• NYU-V2 において、中間層と最終層の両方に蒸留損失を適用したことで、mIoU が 18.75 から 20.75 に向上し、多層蒸留の有効性を示した。
• 本手法はトレーニングダイナミクスを顕著にバランスさせた：蒸留損失はタスク固有損失よりも均等に収束し、SVHN と Omniglot の蒸留損失はそれぞれ 0.089 と 0.041 であったのに対し、タスク固有損失は 0.137 と 1.492 であった。
• NYU-V2 では、セグメンテーション、深度、法線推定の蒸留損失値（0.534、0.381、0.364）が、タスク固有損失（0.027、0.127、0.039）よりもバランスが取れており、最適化のバランスが向上していることを示している。
• 線形アダプタがパフォーマンス向上に十分であり、非線形アダプタが顕著な改善をもたらさなかったため、特徴一致には単純な射影が効果的であると考えられる。
• アダプタを削除すると mIoU が 17.11 に低下し、アダプタによる特徴一致がパフォーマンスとバランスの両面で不可欠であることが確認された。