Compressing audio waveforms using neural networks and vector quantizers

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

receiving an audio waveform that comprises a respective audio sample for each of a plurality of time steps;

processing the audio waveform using an encoder neural network to generate a plurality of feature vectors representing the audio waveform;

generating a respective coded representation of each of the plurality of feature vectors using a plurality of vector quantizers that are each associated with a respective codebook of code vectors, wherein the respective coded representation of each feature vector identifies a plurality of code vectors, including a respective code vector from the codebook of each vector quantizer, that define a quantized representation of the feature vector,

wherein the plurality of vector quantizers are ordered in a sequence, and wherein for each of the plurality of feature vectors, generating the coded representation of the feature vector comprises:

for a first vector quantizer in the sequence of vector quantizers:

receiving the feature vector;

identifying, based on the feature vector, a respective code vector from the codebook of the vector quantizer to represent the feature vector; and

determining a current residual vector based on an error between: (i) the feature vector, and (ii) the code vector that represents the feature vector,

wherein the coded representation of the feature vector identifies the code vector that represents the feature vector; and

generating a compressed representation of the audio waveform by compressing the respective coded representation of each of the plurality of feature vectors.

2. The method of claim 1 , wherein for each of the plurality of feature vectors, generating the coded representation of the feature vector further comprises:

for each vector quantizer after the first vector quantizer in the sequence of vector quantizers:

receiving a current residual vector generated by a preceding vector quantizer in the sequence of vector quantizers;

identifying, based on the current residual vector, a respective code vector from the codebook of the vector quantizer to represent the current residual vector; and

if the vector quantizer is not a last vector quantizer in the sequence of vector quantizers:

updating the current residual vector based on an error between: (i) the current residual vector, and (ii) the code vector that represents the current residual vector;

wherein the coded representation of the feature vector identifies the code vector that represents the current residual vector.

3. The method of claim 1 , wherein generating the compressed representation of the audio waveform comprises:

entropy encoding the respective coded representation of each of the plurality of feature vectors.

4. The method of claim 1 , wherein the respective quantized representation of each feature vector is defined by a sum of the plurality of code vectors identified by the coded representation of the feature vector.

5. The method of claim 1 , wherein the codebooks of the plurality of vector quantizers all include an equal number of code vectors.

6. The method of claim 1 , wherein the encoder neural network and the codebooks of the plurality of vector quantizers are jointly trained along with a decoder neural network, wherein the decoder neural network is configured to:

receive a respective quantized representation of each of a plurality of feature vectors representing an input audio waveform that are generated using the encoder neural network and the plurality of vector quantizers; and

process the quantized representations of the feature vectors representing the input audio waveform to generate an output audio waveform.

7. The method of claim 6 , wherein the training comprises:

obtaining a plurality of training examples that each include: (i) a respective input audio waveform, and (ii) a corresponding target audio waveform;

processing the respective input audio waveform from each training example using the encoder neural network, a plurality of vector quantizers from the sequence of vector quantizers, and the decoder neural network to generate an output audio waveform that is an estimate of the corresponding target audio waveform;

determining gradients of an objective function that depends on the respective output and target waveforms for each training example; and

using the gradients of the objective function to update one or more of: a set of encoder neural network parameters, a set of decoder neural network parameters, or the codebooks of the plurality of vector quantizers.

8. The method of claim 7 , wherein for one or more of the training examples, the target audio waveform is an enhanced version of the input audio waveform.

9. The method of claim 8 , wherein for one or more of the training examples, the target audio waveform is a de-noised version of the input audio waveform.

10. The method of claim 8 , wherein for one or more of the training examples, the target audio waveform is the same as the input audio waveform.

11. The method of claim 10 , wherein processing each input audio waveform to generate the corresponding output audio waveform comprises:

conditioning the encoder neural network, the decoder neural network, or both, on data that defines whether the corresponding target audio waveform is: (i) the input audio waveform or (ii) an enhanced version of the input audio waveform.

12. The method of claim 7 , further comprising, for each training example:

selecting a respective number of vector quantizers to be used in quantizing feature vectors representing the input audio waveform; and

generating the corresponding output audio waveform using only the selected number of vector quantizers from the sequence of vector quantizers.

13. The method of claim 12 , wherein the selected number of vector quantizers to be used in quantizing feature vectors representing input audio waveforms varies between training examples.

14. The method of claim 12 , wherein for each training example, selecting the respective number of vector quantizers to be used in quantizing feature vectors representing the input audio waveform comprises:

randomly sampling the number of vector quantizers to be used in quantizing feature vectors representing the input audio waveform.

15. The method of claim 7 , wherein the objective function comprises a reconstruction loss that, for each training example, measures an error between: (i) the output audio waveform, and (ii) the corresponding target audio waveform.

16. The method of claim 15 , wherein for each training example, the reconstruction loss measures a multi-scale spectral error between: (i) the output audio waveform, and (ii) the corresponding target audio waveform.

17. The method of claim 7 , wherein the training further comprises, for each training example:

processing data derived from the output audio waveform using a discriminator neural network to generate a set of one or more discriminator scores, wherein each discriminator score characterizes an estimated likelihood that the output audio waveform is an audio waveform that was generated using the encoder neural network, the plurality of vector quantizers, and the decoder neural network;

wherein the objective function comprises an adversarial loss that depends on the discriminator scores generated by the discriminator neural network.

18. 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:

receiving an audio waveform that comprises a respective audio sample for each of a plurality of time steps;

processing the audio waveform using an encoder neural network to generate a plurality of feature vectors representing the audio waveform;

generating a respective coded representation of each of the plurality of feature vectors using a plurality of vector quantizers that are each associated with a respective codebook of code vectors, wherein the respective coded representation of each feature vector identifies a plurality of code vectors, including a respective code vector from the codebook of each vector quantizer, that define a quantized representation of the feature vector,

wherein the plurality of vector quantizers are ordered in a sequence, and wherein for each of the plurality of feature vectors, generating the coded representation of the feature vector comprises:

for a first vector quantizer in the sequence of vector quantizers:

receiving the feature vector;

identifying, based on the feature vector, a respective code vector from the codebook of the vector quantizer to represent the feature vector; and

determining a current residual vector based on an error between: (i) the feature vector, and (ii) the code vector that represents the feature vector,

wherein the coded representation of the feature vector identifies the code vector that represents the feature vector; and

generating a compressed representation of the audio waveform by compressing the respective coded representation of each of the plurality of feature vectors.

19. 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:

receiving an audio waveform that comprises a respective audio sample for each of a plurality of time steps;

processing the audio waveform using an encoder neural network to generate a plurality of feature vectors representing the audio waveform;

generating a respective coded representation of each of the plurality of feature vectors using a plurality of vector quantizers that are each associated with a respective codebook of code vectors, wherein the respective coded representation of each feature vector identifies a plurality of code vectors, including a respective code vector from the codebook of each vector quantizer, that define a quantized representation of the feature vector,

wherein the plurality of vector quantizers are ordered in a sequence, and wherein for each of the plurality of feature vectors, generating the coded representation of the feature vector comprises:

for a first vector quantizer in the sequence of vector quantizers:

receiving the feature vector;

identifying, based on the feature vector, a respective code vector from the codebook of the vector quantizer to represent the feature vector; and

determining a current residual vector based on an error between: (i) the feature vector, and (ii) the code vector that represents the feature vector,

wherein the coded representation of the feature vector identifies the code vector that represents the feature vector; and

generating a compressed representation of the audio waveform by compressing the respective coded representation of each of the plurality of feature vectors.

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

receiving an audio waveform that comprises a respective audio sample for each of a plurality of time steps;

processing the audio waveform using an encoder neural network to generate a plurality of feature vectors representing the audio waveform;

generating a respective coded representation of each of the plurality of feature vectors using a plurality of vector quantizers that are each associated with a respective codebook of code vectors, wherein the respective coded representation of each feature vector identifies a plurality of code vectors, including a respective code vector from the codebook of each vector quantizer, that define a quantized representation of the feature vector; and

generating a compressed representation of the audio waveform by compressing the respective coded representation of each of the plurality of feature vectors,

wherein the encoder neural network and the codebooks of the plurality of vector quantizers are jointly trained along with a decoder neural network, and wherein the decoder neural network is configured to:

receive a respective quantized representation of each of a plurality of feature vectors representing an input audio waveform that are generated using the encoder neural network and the plurality of vector quantizers; and

process the quantized representations of the feature vectors representing the input audio waveform to generate an output audio waveform, and

wherein the training comprises:

obtaining a plurality of training examples that each include: (i) a respective input audio waveform, and (ii) a corresponding target audio waveform;

processing the respective input audio waveform from each training example using the encoder neural network, a plurality of vector quantizers from a sequence of vector quantizers, and the decoder neural network to generate an output audio waveform that is an estimate of the corresponding target audio waveform;

determining gradients of an objective function that depends on the respective output and target waveforms for each training example; and

using the gradients of the objective function to update one or more of: a set of encoder neural network parameters, a set of decoder neural network parameters, or the codebooks of the plurality of vector quantizers.