GENERATING PREDICTIONS FOR NON-STATIONARY DATA USING DISTRIBUTIONS OVER OUTPUT HEAD WEIGHTS

1 . A method performed by one or more computers for online inference involving a neural network comprising a base neural network and an output network head, wherein the method comprises:

receiving a data stream comprising a respective input at each of a plurality of time steps; and

processing the data stream to generate a respective predicted output at each time step that estimates a corresponding ground truth output for the time step, comprising, at each time step:

receiving the input at the time step;

obtaining a set of distribution parameters for the time step that parametrizes a transition distribution for the time step,

wherein the transition distribution for the time step defines a conditional probability distribution over possible sets of weights for the time step, given a set of weights for a previous time step;

generating a set of weights for the time step using the set of distribution parameters for the time step;

parametrizing the output network head with the set of weights for the time step; and

processing the input at the time step using the neural network to generate the predicted output at the time step, comprising:

processing the input at the time step using the base neural network to generate an embedding of the input at the time step; and

processing the embedding of the input at the time step using the output network head, in accordance with the set of weights for the time step, to generate the predicted output at the time step.

2 . The method of claim 1 , wherein at each time step:

the transition distribution for the time step is a linear Gaussian transition distribution, and

the set of distribution parameters for the time step comprises a forgetting coefficient for the time step.

3 . The method of claim 2 , wherein at each time step:

the linear Gaussian transition distribution for the time step is variance-preserving, and

the set of distribution parameters for the time step consists of: (i) the forgetting coefficient for the time step, and (ii) a fixed variance.

4 . The method of claim 1 , wherein at each time step, generating the set of weights for the time step using the set of distribution parameters for the time step comprises:

generating the transition distribution for the time step in accordance with the set of distribution parameters for the time step; and

sampling the set of weights for the time step from the transition distribution for the time step.

5 . The method of claim 1 , wherein:

the data stream further comprises the respective ground truth output at each time step, and processing the data stream further comprises, at each time step:

receiving the ground truth output at the time step; and

generating a set of distribution parameters for a next time step using: (i) the set of distribution parameters for the time step, (ii) the embedding of the input at the time step, and (iii) the ground truth output at the time step.

6 . The method of claim 5 , wherein at each time step, generating the set of distribution parameters for the next time step comprises:

generating a predictive posterior distribution for the time step that depends on: (i) the set of distribution parameters for the time step, and (ii) the embedding of the input at the time step,

wherein the predictive posterior distribution for the time step defines a conditional probability distribution over possible ground truth outputs for the time step, given the ground truth output at each previous time step;

determining, from the predictive posterior distribution for the time step, a conditional probability of the ground truth output at the time step, given the ground truth output at each previous time step;

generating an objective function for the time step that depends on the conditional probability of the ground truth output at the time step; and

generating the set of distribution parameters for the next time step by optimizing the objective function for the time step with respect to the set of distribution parameters for the time step.

7 . The method of claim 6 , wherein:

at each time step, the base neural network is configured to process the input at the time step, in accordance with a set of network parameters for the time step, to generate the embedding of the input at the time step, and

processing the data stream further comprises, at each time step:

obtaining the set of network parameters for the time step;

parametrizing the base neural network with the set of network parameters for the time step; and

generating a set of network parameters for the next time step by optimizing the objective function for the time step with respect to the set of network parameters for the time step.

8 . The method of claim 6 , wherein at each time step, the objective function for the time step comprises a logarithm of the conditional probability of the ground truth output at the time step.

9 . The method of claim 6 , wherein at each time step, generating the predictive posterior distribution for the time step comprises:

generating a marginal predictive posterior distribution for the time step that depends on the set of distribution parameters for the time step,

wherein the marginal predictive posterior distribution for the time step defines a conditional probability distribution over possible sets of weights for the time step, given the ground truth output at each previous time step;

generating a likelihood distribution for the time step that depends on the embedding of the input at the time step,

wherein the likelihood distribution for the time step defines a conditional probability distribution over possible ground truth outputs for the time step, given the set of weights for the time step; and

marginalizing, with respect to the possible sets of weights for the time step, the likelihood distribution for the time step over the marginal predictive posterior distribution for the time step.

10 . The method of claim 9 , wherein at each time step, marginalizing, with respect to the possible sets of weights for the time step, the likelihood distribution for the time step over the marginal predictive posterior distribution for the time step comprises:

performing a Monte Carlo estimation over the possible sets of weights for the time step.

11 . The method of claim 9 , wherein at each time step, the likelihood distribution for the time step is a Gaussian likelihood distribution or a softmax likelihood distribution.

12 . The method of claim 9 , wherein at each time step, generating the marginal predictive posterior distribution for the time step comprises:

obtaining a marginal posterior distribution for the previous time step,

wherein the marginal posterior distribution for the previous time step defines a conditional probability distribution over possible sets of weights for the previous time step, given the ground truth output at each previous time step;

generating the transition distribution for the time step in accordance with the set of distribution parameters for the time step; and

marginalizing, with respect to the possible sets of weights for the previous time step, the transition distribution for the time step over the marginal posterior distribution for the previous time step.

13 . The method of claim 9 , wherein processing the data stream further comprises, at each time step:

determining, from the likelihood distribution for the time step, a conditional likelihood of the ground truth output at the time step, given the set of weights for the time step; and

generating, via Bayes' rule, a marginal posterior distribution for the time step in accordance with the conditional likelihood of the ground truth output at the time step and the marginal predictive posterior distribution for the time step,

wherein the marginal posterior distribution for the time step defines a conditional probability distribution over possible sets of weights for the time step, given the ground truth output at the time step and each previous time step.

14 . The method of claim 1 , further comprising, before processing the data stream:

initializing an initial set of distribution parameters that parametrizes a prior distribution,

wherein the prior distribution defines a probability distribution over possible sets of weights prior a first of the plurality of time steps;

generating an initial set of weights using the initial set of distribution parameters; and

initializing the set of distribution parameters for the first time step that parametrizes the transition distribution for the first time step,

wherein the transition distribution for the first time step defines a conditional probability distribution over possible sets of weights for the first time step, give the initial set of weights prior the first time step.

15 . The method of claim 14 , wherein generating the initial set of weights using the initial set of distribution parameters comprises:

generating the prior distribution in accordance with the initial set of distribution parameters; and

sampling the initial set of weights from the prior distribution.

16 . The method of claim 1 , wherein the base neural network is a pre-trained neural network.

17 . The method of claim 1 , wherein the output network head comprises one or more linear neural network layers.

18 . The method of claim 1 , wherein the data stream is non-stationary.

19 . 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 a method for online inference involving a neural network comprising a base neural network and an output network head, wherein the method comprises:

receiving a data stream comprising a respective input at each of a plurality of time steps; and

processing the data stream to generate a respective predicted output at each time step that estimates a corresponding ground truth output for the time step, comprising, at each time step:

receiving the input at the time step;

obtaining a set of distribution parameters for the time step that parametrizes a transition distribution for the time step,

wherein the transition distribution for the time step defines a conditional probability distribution over possible sets of weights for the time step, given a set of weights for a previous time step;

generating a set of weights for the time step using the set of distribution parameters for the time step;

parametrizing the output network head with the set of weights for the time step; and

processing the input at the time step using the neural network to generate the predicted output at the time step, comprising:

processing the input at the time step using the base neural network to generate an embedding of the input at the time step; and

processing the embedding of the input at the time step using the output network head, in accordance with the set of weights for the time step, to generate the predicted output at the time step.

20 . 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 a method for online inference involving a neural network comprising a base neural network and an output network head, wherein the method comprises:

receiving a data stream comprising a respective input at each of a plurality of time steps; and

processing the data stream to generate a respective predicted output at each time step that estimates a corresponding ground truth output for the time step, comprising, at each time step:

receiving the input at the time step;

obtaining a set of distribution parameters for the time step that parametrizes a transition distribution for the time step,

wherein the transition distribution for the time step defines a conditional probability distribution over possible sets of weights for the time step, given a set of weights for a previous time step;

generating a set of weights for the time step using the set of distribution parameters for the time step;

parametrizing the output network head with the set of weights for the time step; and

processing the input at the time step using the neural network to generate the predicted output at the time step, comprising:

processing the input at the time step using the base neural network to generate an embedding of the input at the time step; and

processing the embedding of the input at the time step using the output network head, in accordance with the set of weights for the time step, to generate the predicted output at the time step.