Gradient normalization systems and methods for adaptive loss balancing in deep multitask networks

A method for training a multitask network comprising:

• under control of a hardware processor:
• receiving a training datum of a plurality of training data each associated with a plurality of reference task outputs for the plurality of tasks;
• for each task of the plurality of tasks during a training,
• determining a gradient norm of a single-task loss adjusted by a task weight for the task, with respect to a plurality of network weights of the multitask network, the single-task loss being of (1) a task output for the task determined using a multitask network with the training datum as input, and (2) a corresponding reference task output for the task associated with the training datum;
• and
• determining a relative training rate for the task based on the single-task loss for the task;
• wherein the relative training rates for the plurality of task are determined so that each iteration in the plurality of tasks completes over a similar length of training time;
• determining a gradient loss function comprising a difference between (1) the determined gradient norm for each task and (2) a corresponding target gradient norm determined based on (a) an average gradient norm of the plurality of tasks, and (b) the relative training rate for the task;
• and
• determining an updated task weight for each of the plurality of tasks using a gradient of a gradient loss function with respect to the task weight.

A head mounted display system comprising:

• non-transitory memory configured to store:
• executable instructions, and
• a multitask network;
• a display;
• a sensor;
• and
• a hardware processor in communication with the non-transitory memory and the display, the hardware processor programmed by the executable instructions to:
• determine, using the multitask network, outputs associated with a plurality of tasks, wherein for each task of the plurality of tasks the multitask network is trained by:
• determining a gradient norm of a single-task loss, of (1) a task output for a task of the plurality of tasks determined using the multitask network with a training image as input, and (2) a corresponding reference task output for the task associated with the training image, adjusted by a task weight for the task, with respect to a plurality of network weights of the multitask network,
• determining a relative training rate for the task determined based on the single-task loss for the task,
• determining a gradient loss function comprising a difference between (1) the determined gradient norm for the task and (2) a corresponding target gradient norm determined based on (a) an average gradient norm of the plurality of tasks, (b) the relative training rate for the task, and (c) a hyperparameter, and
• determining an updated task weight for the task using a gradient of the gradient loss function with respect to the task weight for the task such that the updated task weight for the task is based at least in part on the relative training rate for the task;
• wherein the relative training rates for the plurality of tasks are determined so that each iteration in the plurality of tasks completes over a similar length of training time;
• receive a sensor datum captured by the sensor;
• determine a task output for each task of the plurality of tasks using the multitask network with the sensor datum as input;
• and
• cause the display to show information related to the determined task outputs to a user of the head mounted display system.