Training a neural network to perform an algorithmic task using a self-supervised loss

A method performed by one or more computers, the method comprising:

• processing an input dataset using a neural network trained, using a machine learning technique, to perform an algorithmic task;
• wherein the algorithmic task is specified by a computational algorithm defined by a set of rules that, when applied to a dataset, cause the dataset to be processed over a sequence of computational steps to generate an algorithmic output, and
• wherein training the neural network comprises:
• obtaining an input dataset;
• generating a first augmented dataset and a second augmented dataset, wherein for both the first augmented dataset and the second augmented dataset:
• applying the computational algorithm to the augmented dataset causes the same computational operations to be performed at a target computational step as would be performed by applying the computational algorithm to the input dataset;
• processing the first augmented dataset and the second augmented dataset using the neural network, comprising, for each augmented dataset:
• generating an intermediate representation of the augmented dataset at an intermediate layer of the neural network;
• and
• training the neural network on an objective function, wherein the objective function comprises a self-supervised loss term that depends on:
• (i) the intermediate representation of the first augmented dataset generated at the intermediate layer of the neural network, and (ii) the intermediate representation of the second augmented dataset generated at the intermediate layer of the neural network.

A system comprising:

• one or more computers;
• and
• one or more storage devices communicatively coupled to the one or more computers, wherein the one or more storage devices store instructions that, when executed by the one or more computers, cause the one or more computers to perform operations comprising:
• training a neural network to perform an algorithmic task using a machine learning training technique,
• wherein the algorithmic task is specified by a computational algorithm defined by a set of rules that, when applied to a dataset, cause the dataset to be processed over a sequence of computational steps to generate an algorithmic output, and
• wherein training the neural network comprises:
• obtaining an input dataset;
• generating a first augmented dataset and a second augmented dataset, wherein for both the first augmented dataset and the second augmented dataset:
• applying the computational algorithm to the augmented dataset causes the same computational operations to be performed at a target computational step as would be performed by applying the computational algorithm to the input dataset;
• processing the first augmented dataset and the second augmented dataset using the neural network, comprising, for each augmented dataset:
• generating an intermediate representation of the augmented dataset at an intermediate layer of the neural network;
• and
• training the neural network on an objective function, wherein the objective function comprises a self-supervised loss term that depends on:
• (i) the intermediate representation of the first augmented dataset generated at the intermediate layer of the neural network, and (ii) the intermediate representation of the second augmented dataset generated at the intermediate layer of the neural network.

One or more non-transitory computer storage media storing instructions that when executed by one or more computers cause the one or more computers to perform operations comprising:

• training a neural network to perform an algorithmic task using a machine learning training technique,
• wherein the algorithmic task is specified by a computational algorithm defined by a set of rules that, when applied to a dataset, cause the dataset to be processed over a sequence of computational steps to generate an algorithmic output, and
• wherein training the neural network comprises:
• obtaining an input dataset;
• generating a first augmented dataset and a second augmented dataset, wherein for both the first augmented dataset and the second augmented dataset:
• applying the computational algorithm to the augmented dataset causes the same computational operations to be performed at a target computational step as would be performed by applying the computational algorithm to the input dataset;
• processing the first augmented dataset and the second augmented dataset using the neural network, comprising, for each augmented dataset:
• generating an intermediate representation of the augmented dataset at an intermediate layer of the neural network;
• and
• training the neural network on an objective function, wherein the objective function comprises a self-supervised loss term that depends on:
• (i) the intermediate representation of the first augmented dataset generated at the intermediate layer of the neural network, and (ii) the intermediate representation of the second augmented dataset generated at the intermediate layer of the neural network.