Parallel and Flexible Sampling from Autoregressive Models via Langevin Dynamics

Vivek Jayaram, John Thickstun

Introduction

Neural autoregressive models (Larochelle & Murray, 2011) are a popular family of generative models, with wide-ranging applications in a variety of domains including audio (van den Oord et al., 2016a; Dhariwal et al., 2020), images (van den Oord et al., 2016b; Salimans et al., 2017; Parmar et al., 2018; Razavi et al., 2019), and text (Radford et al., 2019; Brown et al., 2020). These models parameterize the conditional distribution over a token in an ordered sequence, given previous tokens in the sequence. The standard approach to sampling from an autoregressive model iteratively generates tokens, according to a conditional distribution over tokens defined by the model, conditioned on the partial sequence of previously generated tokens. We will refer to this approach to sampling as the ancestral sampler.

There are two major drawbacks to ancestral sampling that limit the usefulness of autoregressive models in practical settings. First, ancestral sampling has time complexity that scales linearly in the length of the generated sequence. For data such as high-resolution images or audio, ancestral sampling from an autoregressive model (where the tokens are pixels or sound pressure readings respectively) can be impractically slow. Second, ancestral sampling is frustratingly inflexible. It is easy to sample the second half of a sequence conditioned on the first, but filling in the first half a sequence conditioned on the second naively requires training a new model that reverses the ordering of tokens in the autoregressive factorization. Conditioning on arbitrary subsets of tokens for tasks such as inpainting or super-resolution seems beyond reach of autoregressive modeling.

This paper introduces an alternative, parallel and flexible (PnF) sampler for autoregressive models that can be parallelized and steered using conditioning information or constraints.Code and examples of PnF sampling are available at: https://grail.cs.washington.edu/projects/pnf-sampling/. Instead of sampling tokens sequentially, the PnF sampler initializes a complete sequence (with random tokens) and proceeds to increase the log-likelihood of this sequence by following a Markov chain defined by Langevin dynamics (Neal et al., 2011) on a smoothed log-likelihood. The smoothing temperature is cooled over time according to an annealing schedule informed by Song & Ermon (2019, 2020). Convergence time of this annealed Langevin dynamics is empirically independent of the sequence length and, generalizing Jayaram & Thickstun (2020), the PnF sampler can be applied to posterior log-likelihoods to incorporate conditional information into the sampling process.

The primary technical contribution of this paper is the development of the PnF sampler for discretized autoregressive models (Section 3.1). Our interest in these models is motivated by their success as unconditional models of audio waves (van den Oord et al., 2016a; Mehri et al., 2017; Dhariwal et al., 2020). Defined over a discrete lattice within a continuous space, these models occupy a middle ground between continuous and discrete models. For continuous models such as RNADE (Uria et al., 2013), PnF sampling can be directly applied as in Song & Ermon (2019); Jayaram & Thickstun (2020). We defer the development of the PnF sampler for fully discrete models to future work.

In Section 3.2, we present a stochastic variant of the PnF sampler based on stochastic gradient Langevin dynamics (Welling & Teh, 2011). This is an embarrassingly parallel, asynchronous distributed algorithm for autoregressive sampling. Using a WaveNet model, we show in Section 4.3 that stochastic PnF sampling approximates the quality of ancestral sampling to arbitrary accuracy, with compute time that is inverse proportional to the number of computing devices. This allows PnF sampling to take full advantage of modern, massively parallel computing infrastructure.

We will see in Section 3.3 how the PnF sampler can find solutions to general posterior sampling problems, using an unconditional generative model as a prior. In Section 4 we present applications of the PnF sampler to a variety of Bayesian image and audio inverse problems. We focus on linear inverse problems, using PixelCNN++ (Salimans et al., 2017) and WaveNet (van den Oord et al., 2016a) models as priors. Sections 4.4, 4.5, and 4.6 demonstrate PnF conditional sampling for source separation, super-resolution, inpainting respectively. PnF sampling results correlate strongly with the strength of the generative model used as a prior; as better autoregressive models are developed, they can be used with PnF sampling to improve performance on conditional generation tasks. We refer the reader to the project website for demonstrations of audio PnF sampling.

Related Work

The PnF autoregressive sampler is based on the annealed Langevin dynamics introduced in Song & Ermon (2019), which accelerates standard Langevin dynamics (Neal et al., 2011; Du & Mordatch, 2019) using a smoothing procedure in the spirit of simulated annealing (Kirkpatrick et al., 1983) and graduated optimization (Blake & Zisserman, 1987). The extension of annealed Langevin dynamics to conditional sampling problems was discussed in Jayaram & Thickstun (2020) for source separation and image coloring problems, and developed further in Song et al. (2021) for general posterior sampling problems. The present work extends these methods to discretized autoregressive models, for which the smoothing procedures described in previous work are not directly applicable (Frank & Ilse, 2020). Markov-chain Monte Carlo posterior samplers based on Gibbs sampling rather than Langevin dynamics are proposed in Theis & Bethge (2015) and Hadjeres et al. (2017) as solutions for inpainting tasks.

The slow speed of ancestral sampling is a persistent obstacle to the adoption and deployment of autoregressive models. This has inspired algorithms that seek to parallelize the sampling process. Parallel WaveNet (van den Oord et al., 2018) and ClariNet (Ping et al., 2019) train generative flow models to mimic the behavior of an autoregressive model. Sampling a flow model requires only one pass through a feed-forward network and can be distributed across multiple devices. Wiggers & Hoogeboom (2020) and Song et al. (2020) propose fixed-point algorithms that, like PnF sampling, iteratively refine an initial sample from a simple distribution into a sample from the target distribution. But none of these methods are easily adaptable to source separation (Section 4.4) or more general conditional sampling tasks.

Like anytime sampling (Xu et al., 2021), PnF sampling offers a tradeoff between sample quality and computational budget. The algorithm’s iterates gradually mix to the target distribution and, by stopping early, we can approximate samples from this distribution using less computation. We explore the empirical tradeoff between sample quality and computation for PnF sampling from autoregressive models in Section 4.2. The anytime sampler proposed in Xu et al. (2021) requires a specific model architecture based on the VQ-VAE (van den Oord et al., 2017; Razavi et al., 2019). In contrast, the PnF sampler can be used with any likelihood-based model. Unlike an anytime sampler, the computational budget for PnF sampling must be specified in advance: halting prior to completing the annealing schedule will result in noisy samples.

Bayesian inverse problems are explored extensively in theoretical settings, where the prior is given by a simple analytical distribution (Tropp & Wright, 2010; Knapik et al., 2011; Wang et al., 2017). These problems have also been studied using learned priors given by GAN’s, with a focus on linear inverse problems (Rick Chang et al., 2017; Bora et al., 2017; Raj et al., 2019). These GAN-based approaches are tailored to the latent variable architecture of the model, performing latent space optimizations to find codes that correspond to desired outputs. There is no obvious extension of these latent variable approaches to autoregressive models.

While we focus on autoregressive models, due to their strong empirical performance as unconditional models of audio, PnF sampling could be applied more generally with other likelihood-based models. In the audio space, this includes recent diffusion models (Kong et al., 2021; Chen et al., 2021). Note however that audio vocoder models (Prenger et al., 2019; Kim et al., 2018; Ping et al., 2020), which rely on spectrogram conditioning, cannot be adapted as priors for the source separation, super-resolution, and inpainting experiments presented in Section 4. In addition, GAN based models (Donahue et al., 2019; Kumar et al., 2019), which are not likelihood based, cannot be sampled using PnF.

Parallel and Flexible Sampling

We want to sample from an autoregressive generative model over some indexed sequence of values $\textbf{x}\in\mathcal{X}^{n}$ where

If $p(\textbf{x})$ were a smooth density then, for sufficiently small $\eta$ , the Markov chain mixes and $\textbf{x}^{(t)}$ converges in distribution to $p$ as $t\to\infty$ . But a discretized probability distribution defined over $\mathcal{D}^{n}$ is not smooth; the gradient $\nabla_{\textbf{x}}\log p(\textbf{x}^{(t)})$ is not even well-defined. In Section 3.1 we propose a smoothing of the discrete model $p(\textbf{x})$ , creating a differentiable density on which the Markov chain (2) can mix.

This is a convenient Markov chain for posterior sampling, because the partition function $p(\textbf{y})$ vanishes when we take the gradient. However, like $p(\textbf{x})$ , the posterior $p(\textbf{x}|\textbf{y})$ is not smooth. In Section 3.3 we propose a smoothing of the joint distribution $p(\textbf{x},\textbf{y})$ , for which the posterior $p(\textbf{x}|\textbf{y})$ is differentiable and the Markov chain (3) can mix.

Given a smoothing procedure parameterized by a temperature parameter, we appeal to the simulated annealing heuristic developed in Song & Ermon (2019) to turn down the temperature as the Markov chain (2) or (3) mixes. In contrast to classical Markov chain sampling, for which samples $\textbf{x}^{(t)}$ converge in distribution to $p$ , annealed Langevin dynamics converges asymptotically to a single point distributed approximately according to $p$ . Algorithm 1 describes these annealed Langevin dynamics given a smoothed prior $p_{\sigma}(\textbf{x})$ and smoothed likelihood $p_{\sigma}(\textbf{y}|\textbf{x})$ . The structure of this algorithm is the same as the annealed Langevin dynamics presented in Song & Ermon (2019) (the unconditional case where $\nabla_{\textbf{x}}p_{\sigma}(\textbf{y}|\textbf{x})=0$ ) and Jayaram & Thickstun (2020) (the conditional case); the novel contribution of this paper is the smoothing algorithm for evaluating $\nabla_{\textbf{x}}p_{\sigma}(\textbf{x})$ given $p(\textbf{x})$ (Section 3.1) and $\nabla_{\textbf{x}}p_{\sigma}(\textbf{y}|\textbf{x})$ given $p(\textbf{y}|\textbf{x})$ (Section 3.3).

The functions $f_{i}$ are typically given by a neural network, with shared weights across the sequential indices $i$ . Collectively, these conditional models define the joint distribution $p(\textbf{x})$ according to Equation (1).

2 Stochastic Gradient Langevin Dynamics

This coordinate-wise derivative is only dependent on the tail of the sequence $\textbf{x}_{\geq j}$ :

This doesn’t look promising; calculating an update on a single coordinate $\textbf{x}_{j}$ required $n-j$ inference calculations $p_{\sigma}(\textbf{x}_{i}|\textbf{x}_{<i})$ . But models over very long sequences, including WaveNets, usually make a Markov assumption $p(\textbf{x}_{i}|\textbf{x}_{<i})=p(\textbf{x}_{i}|\textbf{x}_{i-w},\dots,\textbf{x}_{i-1})$ for some limited contextual window of length $w$ . In this case, the coordinate-wise derivative requires only $w$ calls:

Calculating a gradient on a contiguous block of $c$ coordinates leads to a more efficient update

Calculating Equation (12) requires transmission of a block $\{\textbf{x}_{j-w},\dots,\textbf{x}_{j+c+w}\}$ of length $c+2w$ to the computing device, and $c+w$ calculations $p_{\sigma}(\textbf{x}_{i}|\textbf{x}_{<i})$ in order to compute the gradient of a block of length $c$ . If we partition a sequence of length $n$ into $n/c$ blocks of length $c$ , then we can distribute computation of $\nabla_{\textbf{x}}\log p_{\sigma}(\textbf{x})$ with an overhead factor of $1+w/c$ . This motivates choosing $c$ as large as possible, under the constraint that $c+w$ calculations $p_{\sigma}(\textbf{x}_{i}|\textbf{x}_{<i})$ can still be parallelized on a single device.

We can calculate $\nabla_{\textbf{x}}\log p_{\sigma}(\textbf{x})$ by aggregating $n/c$ blocks of gradients according to Equation (12), requiring synchronous communication between $n/c$ machines for every update Equation (2); this is a MapReduce algorithm (Dean & Ghemawat, 2004). We propose a bolder approach in Algorithm 2 based on block-stochastic Langevin dynamics (Welling & Teh, 2011). If $j\in\{1,\dots,n\}$ is chosen uniformly at random then Equation (12) is an unbiased estimate of $\nabla_{\textbf{x}}\log p_{\sigma}(\textbf{x})$ . This motivates block-stochastic updates on patches, which multiple devices can perform asynchronously, a Langevin analog to Hogwild! (Niu et al., 2011).

3 Smoothing a Joint Distribution

We now consider joint distributions $p(\textbf{x},\textbf{y})=p(\textbf{y}|\textbf{x})p(\textbf{x})$ over sources $\textbf{x}\in\mathcal{X}^{n}$ and measurements $\textbf{y}\in\mathcal{Y}^{n}$ . For example, y could be a low resolution version of the signal, or an observed mixture of two signals. We are particularly interested in measurement models of the form $\textbf{y}=g(\textbf{x})$ , for some linear function $g:\mathcal{X}^{n}\to\mathcal{Y}^{n}$ . We can view these measurements y as degenerate likelihoods $p(\textbf{y}|\textbf{x})$ of the form

where $\delta$ denotes the Dirac delta function. This family of linear measurement models describes the source separation, in-painting, and super-resolution tasks featured in Section 4, as well as other linear inverse problems including sparse recovery and image colorization.

For general likelihoods $p(\textbf{y}|\textbf{x})$ (e.g. a classifier) the conditioning values y may depend on the whole sequence x. In this case, stochastic PnF must read the entire sequence x in order to calculate the posterior

But for long sequences x such as audio, the conditioning information y is often a local function of the sequence x. In this case, $\textbf{x}\in\mathcal{X}^{n}$ , $\textbf{y}\in\mathcal{Y}^{n}$ , $\textbf{y}_{i}=g(\textbf{x}_{N(i)})$ where $N(i)$ is a local neighborhood of indices near $i$ , and the likelihood decomposes via conditional independence into

All experiments presented in Section 4 feature this conditioning pattern. For spectrogram conditioning, $N(i)$ is the set of indices (centered at $i$ ) required to compute a short-time Fourier transform. For source separation, super-resolution, and in-painting, $N(i)=i$ . This allows us to compute block gradients of the conditional likelihood (Algorithm 2).

Experiments

We present qualitative and quantitive results of PnF sampling for WaveNet models of audio (van den Oord et al., 2016a) and a PixelCNN++ model of images (Salimans et al., 2017). In Section 4.2 we show that PnF sampling can produce samples of comparable quality to ancestral sampling. In Section 4.3 we show that stochastic PnF sampling is faster than ancestral sampling, when parallelized across a modest number of devices. We go on to demonstrate how PnF sampling can be applied to a variety of image and audio restoration tasks: source separation (Section 4.4), super-resolution (Section 4.5), and inpainting (Section 4.6). We encourage the reader to browse the supplementary material for qualitative examples of PnF audio sampling.

For audio experiments we use the VCTK dataset (Veaux et al., 2016) consisting of 44 hours of speech, as well as the Supra Piano dataset (Shi et al., 2019) consisting of 52 hours of piano recordings. We use a random 80-20 train-test split of VCTK speakers and piano recordings for evaluation. Audio sequences are sampled at a $22$ kHz, with $8$ -bit $\mu$ -law encoding (CCITT, 1988), except for source separation where $8$ -bit linear encoding is used. Sequences used for quantitative evaluation are $50$ k sample excerpts, approximately 2.3 seconds of audio, chosen randomly from the longer test set recordings. For image experiments we use the CIFAR-10 dataset (Krizhevsky, 2009) with the standard train-test split. Additional training and hyperparameter details can be found in the appendix.

2 Quality of Generated Samples

To evaluate the quality of samples generated by PnF sampling, we follow a similar procedure to Holtzman et al. (2020). We compare log-likelihoods, calculated using the noiseless model $p(\textbf{x})$ , of sequences generated by PnF sampling to sequences generated by ancestral sampling from the lowest-noise model $p_{\sigma_{L}}(\textbf{x})$ . Because PnF-sampled sequences are continuous, we quantize these samples to 8-bit values when evaluating their likelihood under the noiseless model. We consider PnF sampling to be successful if it generates sequences with comparable log-likelihoods to ancestral generations.

In Figure 2 we present quantitative results for PnF sampling using an unconditional PixelCNN++ model of CIFAR-10, and a spectrogram-conditioned WaveNet model of both voice and piano datasets. We evaluate $1,000$ generations (length $n=50,000$ sequences for the WaveNet models) using various numbers of Langevin iterations, and report the median log likelihood of quantizations of these sequences under the noiseless model. Asymptotically, as the iterations $T$ of Langevin dynamics increase, the likelihood of PnF samples approaches the likelihood of ancestral samples. Audio PnF samples for various $T$ are presented on the project website.

3 Speed and Parallelism

Ancestral sampling has $O(n)$ serial runtime in the length $n$ of the generated sequence. Using the stochastic PnF sampler described in Section 3.2, the serial runtime is $O(T)$ , where $T$ is the number of Langevin iterations at each level of smoothing. We find empirically that we can set $T$ independent of $n$ , so in principle the serial runtime of stochastic PnF is constant as a function of sequence length. In practice, we do not have an infinite supply of parallel devices, so the serial runtime of stochastic PnF grows inversely proportional to the number of devices. This behavior is demonstrated in Figure 3 for spectrogram-conditioned WaveNet stochastic PnF sampling using a cluster of $8$ Nvidia Titan Xp GPU’s and $T=256$ . Each GPU can calculate Equation (12) for a block of $c=50,000$ samples (2.3 seconds of audio). For PixelCNN++, we find that $T>n$ and therefore the PnF sampler does not improve sampling speed for this model.

Stochastic PnF sampling depends upon asynchronous writes being sparse so that memory overwrites, when two workers update overlapping blocks, are rare. This situation is analogous to the sparse update condition required for Hogwild! If blocks are length $c$ and the number of devices is substantially less than $n/c$ , then updates are sufficiently sparse. But if the number of devices is larger than $n/c$ , memory overwrites become common, and stochastic PnF sampling fails to converge. This imposes a floor on generation time determined by $c$ , exhibited in Figure 3. We cannot substantially reduce this floor by decreasing $c$ because of the tradeoff between $c$ and the model’s Markov window $w$ described in Section 3.2.

In general, PnF sampling becomes faster than ancestral sampling for long sequences $n$ , in which case the stochastic variant of PnF sampling becomes necessary in order to distribute the calculation of $n$ conditional likelihoods. For shorter sequences, accurate unconditional samples can be produced more quickly with the ancestral sampler. Unconditional $32\times 32$ CIFAR-10 generation using the PixelCNN++ mode requires $T=10,000$ Langevin iterations per noise level (Figure 2) for accurate samples; annealing through $L=20$ levels requires a total of $L\times T=200,000$ serial queries to the PixelCNN++ model, far more than $n=3\times 32\times 32$ serial queries to PixelCNN++ for ancestral sampling. Note also that PixelCNN++ conditions on a full image ( $w=n$ ) so the stochastic variant of Pnf sampling is not applicable to this model.

4 Source Separation

We compare PnF audio separation to results using the Demucs (Défossez et al., 2019) and Conv-Tasnet (Luo & Mesgarani, 2019) source separation models. Both Demucs and Conv-Tasnet are supervised models, trained specifically for the source separation task, that learn to output source components given an input mixture. An advantage of PnF sampling is that it does not rely on pairs of source signals and mixes like these supervised methods. We train the supervised models on 10K mixtures of VCTK and Supra Piano samples and measure results on 1K test set mixtures using the standard Scale Invariant Signal-to-Distortion Ratio (SI-SDR) metric for audio source separation (Le Roux et al., 2019). Results in Table 1 show that PnF sampling is competitive with these specialized source separation models. Qualitative comparisons are provided in the supplement. We do not compare results on the popular MusDB dataset (Rafii et al., 2017) because this dataset has insufficient single-channel audio to train WaveNet generative models.

For CIFAR-10, we follow the experimental methodology described in Jayaram & Thickstun (2020). Table 2 shows that PixelCNN++ performs comparably to Glow as a prior, but underperforms NCSN. This makes sense, given the relative strength of NCSN as a prior over CIFAR-10 images in comparison to PixelCNN++ and Glow. Given the strong correlation between the quality of a generative model and the quality of separations using that model as a prior, we anticipate that more recent innovations in autoregressive image models based on transformers (Parmar et al., 2018; Child et al., 2019) will lead to stronger separation results once implementations of these models that match the results reported in these papers become public. Select qualitative image separation results are presented in Figure 4.

5 Super-Resolution

We measure audio super-resolution performance using peak signal-to-noise ratio (PSNR) and compare against a deep learning baseline (Kuleshov et al., 2017) as well as a simple cubic B-spline. Quantitative audio results are presented in Table 3, which show that we outperform these baselines on piano data and produce similar quality reconstructions on voice data. Qualitative audio samples are available in the supplement, where we also show examples of 32x super resolution—beyond the reported ability of existing methods. Select qualitative visual results are presented in Figure 4.

6 Inpainting

Inpainting problems involve the recovery of unobserved data x given a masked observation $g(\textbf{x})=\textbf{m}\odot\textbf{x}$ , where $\textbf{m}\in\{0,1\}^{n}$ (Adler et al., 2011; Pathak et al., 2016). This family of problems includes completion tasks (finishing a sequence given a prime) pre-completion tasks (generating a prefix to a sequence) and outpainting tasks. Ancestral sampling can only be applied to completion tasks, whereas PnF sampling can be used to fill in any pattern of masked occlusions. Qualitative results for audio inpainting are available in the supplement. Select qualitative results for image inpainting are presented in Figure 4.

Conclusion

In this paper we introduced PnF sampling, a parallelizable approach to sampling from autoregressive models that can be flexibly adapted to conditional sampling tasks. The flexibility of PnF sampling decouples the (unconditional) generative modeling problem from the details of specific conditional sampling tasks. Using WaveNet models, we demonstrated a reduction in wall-clock sampling time using PnF sampling in comparison to ancestral sampling, as well as PnF’s ability to solve a variety of practical audio processing problems: source separation, super-resolution, and inpainting. We anticipate that PnF conditional sampling results will improve as developments in generative modeling empower us to incorporate stronger models as priors. More broadly, we are inspired by ongoing research in generative modeling that, coupled with PnF sampling, will continue to drive performance improvements for practical conditional image and audio restoration tasks.

Acknowledgements

We thank Zaid Harchaoui, Sham M. Kakade, Steven Seitz, and Ira Kemelmacher-Shlizerman for valuable discussion and computing resources. This work was supported by a Qualcomm Innovation Fellowship.

References

Appendix A PnF Sampling Details and Hyper-Parameters

We broadly adopt the geometric annealing schedule and hyper-parameters of annealed Langevin dynamics introduced in Song & Ermon (2019) and elaborated upon in Song & Ermon (2020). For both the PixelCNN++ and WaveNet models, we found that we needed additional intermediate noise levels to generate quality samples. We also found that good sample quality using these models required a smaller learning rate and mixing for more iterations than previous work (Song & Ermon, 2019; Jayaram & Thickstun, 2020). We speculate that the need for more levels of annealing and slower mixing could be a attributable to the autoregressive model parameterization, because also required a finer annealing and mixing schedule for the WaveNet models. Detailed hyper-parameters for the PixelCNN++ and WaveNet experiments are presented in Appendix B and C respectively.

Appendix B PixelCNN++ Experimental Details

Our visual sampling experiments are performed using a PixelCNN++ model $p(\textbf{x})$ trained on CIFAR-10. Specifically, we used a public implementation of PixelCNN++ written by Lucas Caccia, available at:

https://github.com/pclucas14/pixel-cnn-pp.

We used the pre-trained weights for this model shared by Lucas Caccia (at the link above) with a reported test-set log loss of 2.95 bits per dimension. For the models $p_{\sigma}(\textbf{x})$ , we fine-tuned the pre-trained model for 10 epochs at each noise level $\sigma^{2}$ . We adopt the geometric annealing schedule proposed in Song & Ermon (2019) for annealing $\sigma^{2}$ , beginning at $\sigma_{1}=1.0$ and ending at $\sigma_{L}=0.01$ using $L=19$ noise levels. This is double the number of noise levels used in previous work (Song & Ermon, 2019; Jayaram & Thickstun, 2020). We also found that sample quality improved using a smaller learning rate and mixing for more iterations than reported in previous work. For conditional sampling tasks, we set $\delta=3e-06$ and $T=300$ in contrast to $\delta=2e-05$ and $T=100$ used in previous work. In wall-clock time, we find that conditional PixelCNN++ sampling tasks require approximately 60 minutes to generate a batch of 16 samples using a 1080Ti GPU.

Appendix C WaveNet Experimental Details

Our audio sampling experiments are performed using a WaveNet model $p(\textbf{x})$ trained on both the VCTK and Supra Piano datasets. We used the public implementation of Wavenet written by Ryuichi Yamamoto available at:

For all audio experiments, where data is encoded between $\{0..255\}$ , we use $L=15$ noise levels geometrically spaced between $\sigma=175.9$ and $\sigma=0.15$ . The same noise levels are also used for the sampling speed and quality results presented in Figures 2 and 3. For all experiments, the number of Langevin steps per noise level is $T=256$ , except for Figure 2 where this parameter is varied to highlight changes in likelihood. The learning rate multiper $\delta=0.05$ is used for all experiments. The Markov window $w$ is based on the underlying architecture. When training the fine-tuned noise models, all training hyperparameters are kept the same as the original WaveNet implementation. For the WaveNet implementation used in this paper, this is $6,139$ samples which is roughly 0.3 seconds at a 22kHz sample rate Please refer to the WaveNet paper or the public WaveNet implementation for training details.

As discussed in the WaveNet paper (van den Oord et al., 2016a), $8$ -bit $\mu$ -law encoding results in a higher fidelity representation of audio than $8$ -bit linear encoding. For most experiments, the observation constraint $y=g(x)$ is still linear even under a $\mu$ -law encoding of $x$ . However, for source separation, the constraint $y=x_{1}+x_{2}$ is no longer linear under $\mu$ -law encoding. Consequently, we use an $8$ -bit linear encoding of $x$ for source separation experiments to avoid a change of variables calculation. To facilitate a fair comparison, all ground truths and baselines shown in the demos use the corresponding $\mu$ -law or linear $8$ -bit encoding.

In the source separation experiments, all mixtures were created with 1/2 gain on each source component. Due to the natural variation of loudness in the training datasets, we find that our model generalizes to mixtures without exactly 1/2 gain on each source. The real life source separation result on the project website shows that we can separate a mixture in the wild when we have no information about the relative loudness of each component.