Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

ATE-SG: Alternate Through the Epochs Stochastic Gradient for Multi-Task Neural Networks

Stefania Bellavia, Francesco Della Santa, Alessandra Papini

TL;DR

This work tackles the problem of conflicting loss gradients in hard-parameter sharing MTNNs by introducing alternate training schemes that separately update shared and task-specific parameters. The proposed SAT-SG and ATE-SG methods enable gradient-based updates that reduce memory and computational costs while maintaining convergence properties akin to classical stochastic gradient descent, with a notable regularization effect during training. The authors provide a rigorous convergence analysis under Lipschitz-gradient assumptions and demonstrate, through synthetic and real-world experiments, that ATE-SG achieves comparable predictive performance to SG while yielding substantial efficiency gains and improved training stability. The approach is designed to be easy to implement in standard deep learning frameworks, and includes an implemented ATE-SG version with an accessible GitHub source for practitioners.

Abstract

This paper introduces novel alternate training procedures for hard-parameter sharing Multi-Task Neural Networks (MTNNs). Traditional MTNN training faces challenges in managing conflicting loss gradients, often yielding sub-optimal performance. The proposed alternate training method updates shared and task-specific weights alternately through the epochs, exploiting the multi-head architecture of the model. This approach reduces computational costs per epoch and memory requirements. Convergence properties similar to those of the classical stochastic gradient method are established. Empirical experiments demonstrate enhanced training regularization and reduced computational demands. In summary, our alternate training procedures offer a promising advancement for the training of hard-parameter sharing MTNNs.

ATE-SG: Alternate Through the Epochs Stochastic Gradient for Multi-Task Neural Networks

TL;DR

This work tackles the problem of conflicting loss gradients in hard-parameter sharing MTNNs by introducing alternate training schemes that separately update shared and task-specific parameters. The proposed SAT-SG and ATE-SG methods enable gradient-based updates that reduce memory and computational costs while maintaining convergence properties akin to classical stochastic gradient descent, with a notable regularization effect during training. The authors provide a rigorous convergence analysis under Lipschitz-gradient assumptions and demonstrate, through synthetic and real-world experiments, that ATE-SG achieves comparable predictive performance to SG while yielding substantial efficiency gains and improved training stability. The approach is designed to be easy to implement in standard deep learning frameworks, and includes an implemented ATE-SG version with an accessible GitHub source for practitioners.

Abstract

This paper introduces novel alternate training procedures for hard-parameter sharing Multi-Task Neural Networks (MTNNs). Traditional MTNN training faces challenges in managing conflicting loss gradients, often yielding sub-optimal performance. The proposed alternate training method updates shared and task-specific weights alternately through the epochs, exploiting the multi-head architecture of the model. This approach reduces computational costs per epoch and memory requirements. Convergence properties similar to those of the classical stochastic gradient method are established. Empirical experiments demonstrate enhanced training regularization and reduced computational demands. In summary, our alternate training procedures offer a promising advancement for the training of hard-parameter sharing MTNNs.
Paper Structure (12 sections, 6 theorems, 29 equations, 7 figures, 6 tables, 3 algorithms)

This paper contains 12 sections, 6 theorems, 29 equations, 7 figures, 6 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Contribution
  3. Multiple-Task Neural Networks
  4. Loss Differentiation and Multiple Tasks
  5. Alternate Training
  6. Simple Alternate Training
  7. Alternate Training through the Epochs
  8. Numerical Experiments
  9. Experiments on Synthetic Data
  10. Experiments on Real-World Data
  11. Conclusion
  12. Implemented ATE-SG

Key Result

Proposition 2.1

Let $\mathrm{N}$ be an MTNN as in def:MTNN and let $\ell,\ell_1,\ldots ,\ell_K$ be the losses in eq:total_loss. Let $\boldsymbol{w}\in\mathbb{R}^p$ be the vector of trainable parameters of $\mathrm{N}$; specifically, $\boldsymbol{w}$ is the concatenation of the shared parameters $\boldsymbol{w}_0\in Then, for each fixed batch $\mathcal{B}$ the vectors are descent directions for the loss $\ell(\ma

Figures (7)

  • Figure 1: Example of MTNN with $2$ tasks. On the left half of the figure, there is $\mathrm{N}_0$ (with input layer in green); on the right half of the figure, there are $\mathrm{N}_1$ and $\mathrm{N}_2$ (with output layers in red).
  • Figure 2: The synthetic dataset $\mathcal{D}$ with respect to its two tasks. Task 1 on the left, task 2 on the right. Different colors denote different labels for the points.
  • Figure 3: Synthetic dataset case.
  • Figure 4: Synthetic dataset case.
  • Figure 5: Synthetic dataset case. Learning rate sequence, starting value: $10^{-3}$. From left to right, from top to bottom: ATE-SG $E_0=E_{\rm ts}=1$, ATE-SG $E_0=E_{\rm ts}=10$, ATE-SG $E_0=E_{\rm ts}=100$, and SG.
  • ...and 2 more figures

Theorems & Definitions (9)

  • Definition 2.1: Hard Parameter Sharing Multi-Task Neural Network
  • Proposition 2.1: Gradients and Descent Directions
  • Lemma 3.1: SAT-SG
  • Theorem 3.2: Robbins-Sigmund 1971 ROBBINS1971
  • Theorem 3.3: SAT-SG convergence
  • Lemma 3.4: ATE-SG
  • Theorem 3.5: ATE-SG convergence
  • Remark 3.1
  • Remark 4.1: MTNN architecture and hyper-parameters