Multi-task gating for machine learning systems

Abstract

Systems and techniques are described herein for training and using multitask machine learning models. For example, a computing device can obtain training data for a first task in a layer in a neural network; perform, based on a determination from a first gating mechanism, the shared function on shared features of the training data using at least one shared channel to generate a shared feature map; perform, based on the determination from the first gating mechanism, the first task-specific function on first task-specific features of the training data using at least one first task-specific channel to generate a first task-specific feature map; generate an output for the first task-specific branch based on performing the shared function on the shared features and performing the first task-specific function on the first task-specific features; and update at least one parameter of the first gating mechanism based on the output.

What is claimed is: 
     
         1 . An apparatus for training a neural network to perform at least one task, comprising:
 at least one memory; and   at least one processor coupled to the at least one memory, the at least one processor configured to:
 obtain training data for a first task in a layer in a neural network, wherein the layer is associated with a first gating mechanism configured to determine whether to process shared features of the training data for the first task using a shared function of a shared branch or first task-specific features of the training data for the first task using a first task-specific function of a first task-specific branch, wherein the shared function is associated with one or more shared channels of the shared branch, and wherein the first task-specific function is associated with one or more first task-specific channels of the first task-specific branch; 
 perform, based on a determination from the first gating mechanism, the shared function on the shared features of the training data for the first task using at least one of the one or more shared channels to generate a shared feature map; 
 perform, based on the determination from the first gating mechanism, the first task-specific function on the first task-specific features of the training data for the first task using at least one of the one or more first task-specific channels to generate a first task-specific feature map; 
 generate an output for the first task-specific branch based on performing the shared function on the shared features of the training data and performing the first task-specific function on the first task-specific features of the training data; and 
 update at least one parameter of the first gating mechanism based on the output. 
   
     
     
         2 . The apparatus of  claim 1 , wherein each shared channel of the one or more shared channels respectively corresponds to each first task-specific channel of the one or more first task-specific channels. 
     
     
         3 . The apparatus of  claim 1 , wherein the first gating mechanism includes a gate for each set of corresponding shared channels and first task-specific channels. 
     
     
         4 . The apparatus of  claim 1 , wherein the at least one parameter of the first gating mechanism includes a plurality of weights. 
     
     
         5 . The apparatus of  claim 4 , wherein the at least one processor is further configured to:
 process the at least one parameter using a sigmoid function to generate a value;   compare the value to a threshold value to provide a binary selection; and   select, based on the binary selection, between the first task-specific function and the shared function.   
     
     
         6 . The apparatus of  claim 1 , wherein the at least one processor is further configured to:
 determine a loss for first task-specific branch based on the output of the first task-specific branch; and   update, using backpropagation, the at least one parameter of the first gating mechanism based on the loss.   
     
     
         7 . The apparatus of  claim 1 , wherein the at least one processor is further configured to:
 obtain training data for a second task in the layer of a second task-specific branch in the neural network, wherein the layer is further associated with a second gating mechanism configured to determine whether to process shared features of the training data for the second task using the shared function of the shared branch or second task-specific features of the training data for the second task using a second task-specific function, and wherein the second task-specific function is associated with one or more second task-specific channels of the second task-specific branch;   perform, based on a determination from the second gating mechanism, the shared function on the shared features of the training data for the second task using at least one of the one or more shared channels of the shared feature map to generate a second shared feature map;   perform, based on a determination from the second gating mechanism, the second task-specific function on the second task-specific features of the training data for the second task using at least one of the one or more second task-specific channels to generate a second task-specific feature map;   generate a second output for the second task-specific branch based on performing the shared function on the shared features of the training data and performing the second task-specific function on the second task-specific features; and   update at least one parameter of the second gating mechanism based on the second output.   
     
     
         8 . The apparatus of  claim 7 , wherein the at least one processor is further configured to:
 determine a second loss for the second task-specific branch based on the output of the second task-specific branch; and   update, using backpropagation, the at least one parameter of the second gating mechanism based on the second loss.   
     
     
         9 . The apparatus of  claim 7 , wherein a first subset of the one or more shared channels is allocated to a first subset of the training data for the first task and a second subset of the one or more shared channels is allocated to the first subset of the training data for the second task. 
     
     
         10 . The apparatus of  claim 7 , wherein a classification of the first task is different from a classification of the second task. 
     
     
         11 . The apparatus of  claim 1 , wherein the at least one processor is further configured to:
 based on the first gating mechanism, select between at least one channel of the one or more first task-specific channels and at least one corresponding channel of the one or more shared channels; and   prune unselected channels based on the selection by the first gating mechanism.   
     
     
         12 . The apparatus of  claim 1 , wherein a number of active first task-specific channels is different from a number of active shared channels. 
     
     
         13 . The apparatus of  claim 1 , wherein a first channel of the one or more shared channels corresponds to a first channel of the one or more first task-specific channels. 
     
     
         14 . The apparatus of  claim 13 , wherein the first channel of the one or more shared channels is active and the first channel of the one or more first task-specific channels is inactive. 
     
     
         15 . The apparatus of  claim 13 , wherein the first channel of the one or more shared channels in inactive and the first channel of the one or more first task-specific channels is active. 
     
     
         16 . The apparatus of  claim 1 , wherein the first task is one of image segmentation, surface normal estimation, depth estimation, or classification. 
     
     
         17 . An apparatus for performing at least one task, comprising:
 at least one memory; and   at least one processor coupled to the at least one memory, the at least one processor configured to:
 receive input data for a first task in a layer in a neural network, wherein the layer is associated with a shared function of a shared branch and a first task-specific function of a first task-specific branch, and wherein the shared function is associated with one or more shared channels of the shared branch, and wherein the first task-specific function is associated with one or more first task-specific channels of the first task-specific branch; 
 perform the shared function on shared features of the input data for the first task using at least one of the one or more shared channels of the shared branch to generate a shared feature map; 
 perform the first task-specific function on first task-specific features of the input data using at least one of the one or more first task-specific channels associated with the first task-specific function to generate a first task-specific feature map; and 
 generate an output for the first task-specific branch based on performing the shared function on the shared features of the input data and performing the first task-specific function on the first task-specific features of the input data. 
   
     
     
         18 . The apparatus of  claim 17 , wherein the at least one processor is further configured to:
 receive input data for a second task in the layer in the neural network, wherein the layer is further associated with a second task-specific function, and wherein the second task-specific function is associated with one or more second task-specific channels;   perform the shared function on shared features of the input data for the second task using at least one of the one or more shared channels of the shared feature map to generate a second shared feature map;   perform the second task-specific function on second task-specific features of the input data for the second task using at least one of the one or more second task-specific channels to generate a second task-specific feature map; and   generate an output for the first task-specific branch based on performing the shared function on the shared features of the input data and performing the first task-specific function on the first task-specific features of the input data.   
     
     
         19 . The apparatus of  claim 17 , wherein the at least one processor is further configured to:
 receive the output in a subsequent layer of the neural network, wherein the subsequent layer includes a third task-specific function of the first task-specific branch and a second shared function of the shared branch, and wherein the third task-specific function includes a plurality of third task-specific channels, and wherein the second shared function includes a plurality of shared channels;   perform the third task-specific function on third task-specific features of the output using at least one of the plurality of third task-specific channels to generate a first subsequent feature map;   perform the second shared function on shared features of the output using at least one of the plurality of shared channels to generate a second subsequent feature map; and   generate a subsequent output for the first task-specific branch based on performing the second shared function on the shared features of the output and performing the third task-specific function on the third task-specific features of the output.   
     
     
         20 . A processor-implemented method for training a neural network to perform at least one task, comprising:
 obtain training data for a first task in a layer in a neural network, wherein the layer is associated with a first gating mechanism configured to determine whether to process shared features of the training data for the first task using a shared function of a shared branch or first task-specific features of the training data for the first task using a first task-specific function of a first task-specific branch, wherein the shared function is associated with one or more shared channels of the shared branch, and wherein the first task-specific function is associated with one or more first task-specific channels of the first task-specific branch;   perform, based on a determination from the first gating mechanism, the shared function on the shared features of the training data for the first task using at least one of the one or more shared channels to generate a shared feature map;   perform, based on the determination from the first gating mechanism, the first task-specific function on the first task-specific features of the training data for the first task using at least one of the one or more first task-specific channels to generate a first task-specific feature map;   generate an output for the first task-specific branch based on performing the shared function on the shared features of the training data and performing the first task-specific function on the first task-specific features of the training data; and   update at least one parameter of the first gating mechanism based on the output.   
     
     
         21 . The processor-implemented method of  claim 20 , wherein each shared channel of the one or more shared channels respectively corresponds to each first task-specific channel of the one or more first task-specific channels. 
     
     
         22 . The processor-implemented method of  claim 20 , wherein the first gating mechanism includes a gate for each set of corresponding shared channels and first task-specific channels. 
     
     
         23 . The processor-implemented method of  claim 20 , wherein the at least one parameter of the first gating mechanism includes a plurality of weights. 
     
     
         24 . The processor-implemented method of  claim 23 , further comprising:
 processing the at least one parameter using a sigmoid function to generate a value;   comparing the value to a threshold value to provide a binary selection; and   selecting, based on the binary selection, between the first task-specific function and the shared function.   
     
     
         25 . The processor-implemented method of  claim 20 , further comprising:
 determining a loss for first task-specific branch based on the output of the first task-specific branch; and   updating, using backpropagation, the at least one parameter of the first gating mechanism based on the loss.   
     
     
         26 . The processor-implemented method of  claim 20 , further comprising:
 obtaining training data for a second task in the layer of a second task-specific branch in the neural network, wherein the layer is further associated with a second gating mechanism configured to determine whether to process shared features of the training data for the second task using the shared function of the shared branch or second task-specific features of the training data for the second task using a second task-specific function, and wherein the second task-specific function is associated with one or more second task-specific channels of the second task-specific branch;   performing, based on a determination from the second gating mechanism, the shared function on the shared features of the training data for the second task using at least one of the one or more shared channels of the shared feature map to generate a second shared feature map;   performing, based on a determination from the second gating mechanism, the second task-specific function on the second task-specific features of the training data for the second task using at least one of the one or more second task-specific channels to generate a second task-specific feature map;   generating a second output for the second task-specific branch based on performing the shared function on the shared features of the training data and performing the second task-specific function on the second task-specific features; and   updating at least one parameter of the second gating mechanism based on the second output.   
     
     
         27 . The processor-implemented method of  claim 26 , further comprising:
 determining a second loss for the second task-specific branch based on the output of the second task-specific branch; and   updating, using backpropagation, the at least one parameter of the second gating mechanism based on the second loss.   
     
     
         28 . The processor-implemented method of  claim 26 , wherein a first subset of the one or more shared channels is allocated to a first subset of the training data for the first task and a second subset of the one or more shared channels is allocated to the first subset of the training data for the second task. 
     
     
         29 . The processor-implemented method of  claim 26 , wherein a classification of the first task is different from a classification of the second task. 
     
     
         30 . The processor-implemented method of  claim 20 , further comprising:
 based on the first gating mechanism, selecting between at least one channel of the one or more first task-specific channels and at least one corresponding channel of the one or more shared channels; and   pruning unselected channels based on the selection by the first gating mechanism.