PaperlayerView exact revision on arXiv

Source-linked AI summary

It's Raw! Audio Generation with State-Space Models

Karan Goel, Albert Gu, Chris Donahue, Christopher Ré

arXiv:2202.09729v1cs.SDcs.AIcs.LGeess.AS

TL;DR

Raw-audio generation must model high-rate waveforms with long-range structure while avoiding unfavorable computational tradeoffs. The paper introduces SaShiMi, a stabilized S4-based multiscale architecture, and reports strong autoregressive and non-autoregressive results. SaShiMi also improves diffusion generation when replacing a WaveNet backbone.

  • Problem

    Raw-audio generation is challenging because waveforms have high sampling rates and long-range behavior across multiple timescales.

  • Method

    SaShiMi combines stabilized S4 layers with pooling-based multiscale architecture and bidirectional relaxation for non-autoregressive models.

  • Results

    SaShiMi achieves stronger autoregressive waveform-generation results and improves diffusion generation when used as a WaveNet backbone replacement.

  • Takeaways & Limitations

    SSM-based architectures offer a strong alternative to conventional audio architectures across training speed, generation speed, sample efficiency, and audio quality.

  • Takeaways & Limitations

    The work focuses primarily on unconditional raw-waveform generation, while much prior waveform generation uses conditioning information.

Abstract

from arXiv · show

Developing architectures suitable for modeling raw audio is a challenging problem due to the high sampling rates of audio waveforms. Standard sequence modeling approaches like RNNs and CNNs have previously been tailored to fit the demands of audio, but the resultant architectures make undesirable computational tradeoffs and struggle to model waveforms effectively. We propose SaShiMi, a new multi-scale architecture for waveform modeling built around the recently introduced S4 model for long sequence modeling. We identify that S4 can be unstable during autoregressive generation, and provide a simple improvement to its parameterization by drawing connections to Hurwitz matrices. SaShiMi yields state-of-the-art performance for unconditional waveform generation in the autoregressive setting. Additionally, SaShiMi improves non-autoregressive generation performance when used as the backbone architecture for a diffusion model. Compared to prior architectures in the autoregressive generation setting, SaShiMi generates piano and speech waveforms which humans find more musical and coherent respectively, e.g. 2x better mean opinion scores than WaveNet on an unconditional speech generation task. On a music generation task, SaShiMi outperforms WaveNet on density estimation and speed at both training and inference even when using 3x fewer parameters. Code can be found at https://github.com/HazyResearch/state-spaces and samples at https://hazyresearch.stanford.edu/sashimi-examples.

1 Introduction

Raw-audio generation requires models that handle high-rate waveforms, long-range dependencies, computational efficiency, and sample efficiency. SaShiMi combines stabilized S4 layers with a multiscale architecture and achieves strong autoregressive and non-autoregressive generation results.

  • Raw-audio waveforms contain tens of thousands of timesteps per second and exhibit long-range behavior across multiple timescales.
  • Existing audio sequence models face tradeoffs: RNNs cannot parallelize training, while CNNs require parameters proportional to context length.
  • S4 is reparameterized using state-space stability theory so its recurrent representation remains stable during autoregressive generation.
  • SaShiMi combines S4 layers with pooling, residual connections, and multiscale transformations to capture hierarchical information across resolutions.
  • SaShiMi supports both autoregressive and non-autoregressive generation through a causal S4 design and a bidirectional relaxation for global context.
  • 128K vs 4K context lengths and more than 0.15 bits better negative log-likelihoods distinguish SaShiMi from WaveNet in long-context music modeling.
  • SaShiMi improves unconditional speech generation and diffusion-based non-autoregressive generation, while offering favorable speed, sample-efficiency, and audio-quality tradeoffs.

2 Related Work

Prior waveform-generation systems use autoregressive or non-autoregressive approaches, with WaveNet serving as a widely reused architecture. SaShiMi is presented as an alternative that can leverage unbounded context and transfer to non-autoregressive settings.

  • Unconditional waveform generation is harder than conditioned generation because one model must learn both local and global structure.
  • Autoregressive models generate samples sequentially and offer tractable exact likelihoods, whereas non-autoregressive models generate entire waveforms in one pass.
  • WaveNet is widely used in vocoding, flow models, autoencoders, GANs, and non-autoregressive audio generation systems.
  • WaveNet cannot model long-term structure beyond its receptive field, motivating SaShiMi as an alternative that can leverage unbounded context.
  • SaShiMi transfers to non-autoregressive settings in addition to its primary autoregressive evaluation.

3 Background

Autoregressive waveform modeling factorizes audio likelihood into next-sample predictions, but CNNs and RNNs trade off training and inference efficiency while struggling with long contexts. State-space models combine convolutional and recurrent views to support efficient long-range modeling.

  • 3.1 Autoregressive Modeling of Audio: Autoregressive models factorize waveform likelihood into conditional next-sample distributions, enabling tractable exact likelihood training and sequential generation.Training maximizes likelihood over a sequence; inference samples each next waveform value from its preceding context.
  • 3.1 Autoregressive Modeling of Audio: Dilated convolutions enlarge CNN receptive fields, but modeling context size T with CNNs requires parameters proportional to T.This parameter growth is problematic for audio sequences requiring very large contexts.
  • 3.1 Autoregressive Modeling of Audio: CNNs parallelize training but require specialized inference caching, whereas RNNs infer efficiently but train sequentially and face optimization difficulties.CNN context size also grows with kernel widths, while RNN hidden states summarize the entire preceding sequence.
  • 3.2 State Space Models: State-space models provide linear recurrent and convolutional forms, combining efficient training and inference while avoiding RNN optimization and CNN context-size limitations.Their convolutional representation has an unbounded kernel, while their recurrent representation uses linear state updates.
  • 3.3 S4: S4 uses a diagonal-plus-low-rank state matrix that enables fast convolution-kernel computation and incorporates HiPPO matrices for long-range dependencies.S4 initializes its state matrices using the HiPPO formulation.

4 Model

SaShiMi builds waveform models from modified S4 layers arranged across multiple resolutions. Its tied parameterization targets recurrent-generation stability, while pooling connects residual S4 stacks across raw and downsampled audio resolutions.

  • 4 Model: SaShiMi combines modified S4 layers for long-context efficiency with a multi-scale architecture connecting stacks of residual S4 blocks.The top tier processes raw-rate audio, while lower tiers process downsampled inputs whose outputs are upsampled and combined.
  • 4.1 Stabilizing S4 for Recurrence: SaShiMi replaces S4’s Λ + pq∗ parameterization with Λ −pp∗, effectively tying p and q while reversing a sign.The restricted form remains a diagonal-plus-low-rank matrix and retains S4’s fast computation and HiPPO initialization.
  • 4.1 Stabilizing S4 for Recurrence: A Hurwitz matrix has eigenvalues with negative real parts, ensuring asymptotic stability for the state-space model.The paper uses the negative-semidefinite term −pp∗ to make stability easier to establish and enforce.
  • 4.1 Stabilizing S4 for Recurrence: S4’s recurrent mode requires a Hurwitz state matrix, but its convolutional training mode does not require this condition because it avoids repeatedly powering A.The modified parameterization makes the spectrum easier to control through the diagonal portion Λ.
  • 4.2 SaShiMi Architecture: Pooling layers consolidate waveform information across resolutions, with experiments using two pooling layers and three tiers.Pooling reshapes and linearly transforms sequences while increasing hidden dimension through an expansion factor.
  • 4.2 SaShiMi Architecture: Bidirectional S4 processes the sequence in both forward and reverse directions, then concatenates the outputs through a positionwise linear layer.The paper reports that this variant outperforms causal S4 when autoregression is not required.

5 Experiments

Experiments evaluate SaShiMi across music and speech generation in autoregressive and non-autoregressive settings. SaShiMi improves likelihood, musicality, computational efficiency, stability, and speech-generation quality relative to baselines.

  • Experimental setup: SaShiMi is evaluated on music and speech datasets in both autoregressive and non-autoregressive settings.The benchmarks include Beethoven, YouTubeMix, and SC09, with comparisons against SampleRNN, WaveNet, and DiffWave.
  • Unbounded music generation: 0.40 points: SaShiMi improves musicality MOS on 16-second YouTubeMix samples while achieving similar audio fidelity to other methods.The result supports more coherent generation beyond the training context.
  • Architectural ablations: SaShiMi trains more stably and efficiently than baselines in wall-clock time, while pooling improves computation and modeling over isotropic stacked S4 layers.The architectural ablation controls for parameter count when comparing pooled and isotropic designs.
  • S4 stability: The Hurwitz parameterization allows A to be learned while preserving generation stability, addressing instability observed with standard S4 parameterization.Figure 3 compares the spectral radii of A matrices under different parameterizations.
  • Efficiency tradeoffs: A model more than 3X smaller simultaneously outperforms all baselines on NLL and computation at both training and inference.Performance improves monotonically with depth, although additional depth increases computation.
  • Unconditional speech generation: 2× higher MOS: SaShiMi improves speech quality and intelligibility over other autoregressive waveform models on SC09.Annotators also agree on SaShiMi sample labels more often, while the DiffWave integration improves all metrics without tuning.

6 Discussion

SaShiMi produces more musical and intelligible waveforms than previous architectures while improving training and inference efficiency and sample efficiency. Its S4-based design also transfers to non-autoregressive generation, where replacing WaveNet improves DiffWave.

  • Discussion: Human evaluations judge SaShiMi waveforms more musical for music and more intelligible for speech than waveforms from previous architectures.The paper interprets these outcomes as evidence of greater global coherence.
  • Discussion: SaShiMi is more computationally efficient than past architectures during both training and inference by leveraging S4’s convolutional and recurrent forms.The discussion connects this efficiency to the underlying S4 representation.
  • Discussion: SaShiMi achieves better quantitative performance with fewer training steps, indicating consistent sample-efficiency gains across the evaluated settings.The claim is stated alongside improvements in audio quality and computational efficiency.
  • Discussion: Using SaShiMi as a drop-in WaveNet replacement improves an existing state-of-the-art unconditional generation model without additional tuning.This result supports its use as a backbone beyond autoregressive waveform modeling.

A.2 Model Architecture

The S4 block combines an S4 layer with position-wise neural-network operations. Its feed-forward projections expand and then contract the feature dimension.

  • The first portion of the S4 block follows the S4 block used by Gu et al.
  • The block applies LayerNorm, S4, a GELU activation, and a linear projection in sequence.GELU is the activation used in the implementation.
  • All operations except the S4 layer are position-wise along the time or sequence dimension.
  • The feed-forward projections use W1 ∈ R^d×e^d and W2 ∈ R^e^d×d with expansion factor e = 2.

B Additional Results

The additional-results section reports architectural ablations and efficiency benchmarking for the evaluated baselines.

  • The paper provides details of architectural ablations.
  • The paper provides details of efficiency benchmarking.
  • The reported ablations and benchmarks complement the main experimental results.

B.0.1 YouTubeMix

The YouTubeMix experiments compare multiple architectural variants and benchmark throughput across batch sizes. SaShiMi is evaluated against SampleRNN, WaveNet, and isotropic S4 variants.

  • YouTubeMix: The YouTubeMix evaluation conducts architectural ablations and efficiency benchmarking for all baselines.
  • Efficiency benchmarking: SaShiMi throughput scales near linearly with batch size, unlike WaveNet’s sublinear scaling and SampleRNN’s earlier throughput peak.
  • Efficiency benchmarking: 3× and 5×: SaShiMi-2 improves peak throughput over WaveNet and SampleRNN, respectively.
  • Architectures: The ablations vary the number of S4 blocks per tier across SaShiMi-{2, 4, 6, 8} and compare SampleRNN and WaveNet variants.
  • Architectures: Isotropic S4 removes pooling from SaShiMi, but its larger memory usage limits tested sequence lengths to 4s or 2s for the 4- and 8-layer models.
  • Efficiency benchmarking: Peak throughput is measured by generating 1000 samples across batch sizes from 1 to 8192 in powers of 2.
  • Diffusion model ablations: Diffusion-model ablation results are reported in Table 7, with experimental details in Appendix C.2.

C.1 Datasets

The experiments use 16kHz audio datasets with dataset-specific preprocessing, quantization, model configurations, and training schedules. Autoregressive and diffusion models are trained on different hardware settings.

  • Dataset preparation: All datasets are preprocessed to 16kHz using torchaudio.
  • Dataset preparation: Beethoven experiments use linear quantization, while other datasets use quantization settings based on prior work.
  • Datasets: YouTubeMix contains 4 hours of piano music split into 1-minute files, with 88% for training, 6% validation, and 6% testing.
  • Datasets: SC09 uses 31,158 one-second spoken-digit utterances from 2,032 speakers, totaling 8.7 hours at 16kHz.
  • Training setup: SaShiMi, SampleRNN, and WaveNet receive 8-bit quantized inputs without additional data augmentation.
  • Autoregressive models: SampleRNN uses truncated backpropagation through time with a chunk size of 1024, whereas WaveNet and SaShiMi train on entire chunks.
  • Autoregressive models: SampleRNN can improve steadily before suddenly diverging, indicating instability during training.
  • Diffusion models: The large SaShiMi diffusion model uses hidden dimension 128, six S4 blocks per tier, two pooling layers, and pooling factor 4.

C.3 Automated Evaluations

The paper evaluates autoregressive models using likelihood-based and classifier-based metrics, supplemented by human-agreement measurements. These procedures assess both predictive fit and generated-sample quality, diversity, and consistency with human judgments.

  • Likelihood evaluation: Negative log-likelihood is reported in bits on each dataset’s test set using non-overlapping chunks matching training length.Predictions at every step of each chunk are used to calculate the chunk’s average NLL.
  • Generated-sample evaluation: Generated samples are evaluated with four standard metrics using a ResNeXT classifier achieving 98.3% test accuracy.The reported metrics are FID, IS, mIS, and AM Score; NDB is omitted because it was slow and generally uninformative.
  • Generated-sample evaluation: FID compares generated and real samples in classifier feature space, while IS measures classifier confidence alongside sample quality and diversity.These metrics provide complementary assessments of generated audio distributions.
  • Generated-sample evaluation: mIS measures intra-class and inter-class diversity, whereas AM Score compares the generated marginal label distribution with that of the training data.Both metrics extend evaluation beyond classifier confidence alone.
  • Human agreement: Cohen’s inter-annotator agreement κ treats the classifier and aggregated crowdworker digit predictions as separate raters.The crowdworkers are treated as a single rater for this calculation.

C.3.1 Evaluation Procedure for Autoregressive Models

The evaluation procedure combines likelihood-filtered automated metrics with crowdsourced judgments of generated audio. Separate protocols assess unconditional music and speech generation, controlling sample selection, presentation, and human-rating tasks.

  • Automated evaluation: 5120 samples per autoregressive model are ranked by likelihood, with the lowest 0.40 and highest 0.05 fractions discarded before automated evaluation.The remaining samples are used for metric calculation, and the cutoff thresholds are validated on a separate set of 5120 samples.
  • Automated evaluation: Automated metrics are computed on 2048 random samples generated from each model.
  • YouTubeMix evaluation: YouTubeMix MOS evaluates audio fidelity and musicality for unconditional SaShiMi, SampleRNN, and WaveNet samples.The protocol uses 1024 generated 16-second samples per method, with direct autoregressive sampling.
  • YouTubeMix evaluation: For YouTubeMix, noise-like or mostly silent samples are rejected, 25 dataset clips and gold-standard clips are added, and workers rate randomized method batches.The study collects 20 responses per batch through Amazon Mechanical Turk, paying $0.50 per batch.
  • Crowdsourcing interfaces: YouTubeMix uses a crowdsourcing interface presenting one sample from each method and the dataset, whereas SC09 presents 10 samples from one method per batch.The interfaces collect the task-specific fidelity, musicality, intelligibility, quality, and diversity judgments.
  • SC09 evaluation: SC09 MOS measures digit intelligibility, audio quality, and speaker diversity, while workers also classify spoken digits.Protocols cover autoregressive, WaveGAN, and diffusion models, with diffusion samples generated using 200 denoising steps.
  • SC09 evaluation: For SC09, a ResNeXT classifier selects the top 50 samples within each digit category, which are grouped into 50 batches per method or dataset.Each batch contains 10 randomized samples, and workers provide per-file digit judgments plus batch-level quality and diversity ratings.
Loading 2202.09729v1…