Flexible Parameter Sharing for Multi-Task Learning

A computing system, comprising:

• at least one processor;
• a multi-task machine-learned model configured to perform a plurality of tasks T, comprising:
• a plurality of layers L, each layer comprising a plurality of components C;
• a routing matrix of size T×C associated with each respective layer, the routing matrix for a particular layer comprising a matrix of binary allocation variables descriptive of which components in the respective layer an input into the machine-learned model is routed through to generate an output;
• and
• a plurality of task-specific heads, each task-specific head configured to receive an output from a final layer of the one or more layers and generate an output associated with a respective task;
• and
• at least one tangible, non-transitory computer-readable medium that stores instructions that, when executed by the at least one processor, cause the at least one processor to perform operations, the operations comprising:
• obtaining an input;
• selecting a particular task;
• routing the input through the machine-learned model according to the respective routing matrix for each layer for the particular task;
• and
• receiving, as an output of the machine-learned model, a task-specific output from the task-specific head associated with the particular task;
• wherein the multi-task machine-learned model has been trained using a straight-through Gumbel-softmax approximation to jointly learn the routing matrix with the plurality of components using back-propagation.

One or more tangible, non-transitory computer-readable media that store instructions that, when executed by one or more processors, cause the one or more processors to perform operations, the operations comprising:

• obtain a test input for a machine-learned model configured to perform a plurality of tasks, the machine-learned model comprising a plurality of layers, each layer comprising a plurality of components, each task assigned to select one or more components for each layer according to a connection probability matrix for each respective layer comprising a matrix of connection probabilities for each component to be used in the respective layer for the task, each connection probability comprising two complementary logits;
• selecting a particular task from the one or more tasks;
• and
• training the machine-learned model for the particular task, wherein training the machine-learned model for the particular task comprises:
• performing a forward pass using the test input and the connection probability matrix for each layer to generate a sample distribution of test outputs;
• training the components of the machine-learned model based at least in part on the sample distribution;
• and
• performing a backwards pass to train the connection probability matrix of the machine-learned model using a straight-through Gumbel-softmax approximation.