Enhancing into the codec: Noise Robust Speech Coding with Vector-Quantized Autoencoders
Audio codecs based on discretized neural autoencoders have recently been developed and shown to provide significantly higher compression levels for comparable quality speech output. However, these models are tightly coupled with speech content, and produce unintended outputs in noisy conditions. Based on VQ-VAE autoencoders with WaveRNN decoders, we develop compressor-enhancer encoders and accompanying decoders, and show that they operate well in noisy conditions. We also observe that a compressor-enhancer model performs better on clean speech inputs than a compressor model trained only on clean speech.
main.tex
% Template for ICASSP-2021 paper; to be used with:
%          spconf.sty  - ICASSP/ICIP LaTeX style file, and
%          IEEEbib.bst - IEEE bibliography style file.
% --------------------------------------------------------------------------
\documentclass{article}
\usepackage{spconf,amsmath,graphicx}
\usepackage[ruled, vlined, linesnumbered]{algorithm2e}
\usepackage{bbm}
\usepackage{float}
\usepackage{multicol}
\usepackage{makecell}
\DeclareMathOperator*{\argmin}{arg\,min}
% Title.
% ------
\title{Enhancing into the codec: Noise Robust Speech Coding with Vector-Quantized Autoencoders}

\name{Jonah Casebeer\sthanks{Equal contribution. Work performed while at Amazon Web Services.}$^{\dagger}$ \qquad Vinjai Vale\footnotemark[1]$^{\flat}$  \qquad  Umut Isik$^{\sharp}$ \qquad}

\secondnm{\textit{Jean-Marc Valin}$^{\sharp}$ \qquad \textit{Ritwik Giri}$^{\sharp}$ \qquad \textit{Arvindh Krishnaswamy}$^{\sharp}$}

\address{$^{\dagger}$ University of Illinois at Urbana-Champaign, $^{\flat}$ Stanford University\\
          $^{\sharp}$ Amazon Web Services
}
         

\begin{document}
%\ninept
%
\maketitle
%

\begin{abstract}
Audio codecs based on discretized neural autoencoders have recently been developed and shown to provide significantly higher compression levels for comparable quality speech output. However, these models are tightly coupled with speech content, and produce unintended outputs in noisy conditions. Based on VQ-VAE autoencoders with WaveRNN decoders, we develop compressor-enhancer encoders and accompanying decoders, and show that they operate well in noisy conditions. We also observe that a compressor-enhancer model performs better on clean speech inputs than a compressor model trained only on clean speech. 
\end{abstract}
%

%
\begin{keywords}
speech enhancement, speech coding, audio compression
\end{keywords}
%


\input{sections/intro.tex}
\input{sections/method.tex}
\input{sections/exp_results.tex}
\input{sections/conclusion.tex}


% References should be produced using the bibtex program from suitable
% BiBTeX files (here: strings, refs, manuals). The IEEEbib.bst bibliography
% style file from IEEE produces unsorted bibliography list.
% -------------------------------------------------------------------------
\bibliographystyle{IEEEbib}
\bibliography{strings,refs}

\end{document}

% -- switch to vctk only no dns DONE
% -- 512 speaker embeddings and more layers in speaker embedding net
% -- frame vqvae + wavenet not in bad light
% -- experiments table 1 make clear relationship 
% -- maybe remove table 2

%%%%%%%%%%%%%%%%% jm valin %%%%%%%%%%%%%%%%%%%%%
% be more clear on what a compressor enhanceer is -- TRIED
% make more clear the difference between this and VQVAE wavenet combo
% unclear if real-time or not -- say it is not now DONE
% make sure to define bit-rates and talk about what they are -- have added in several and fleshed out lit review
% 

%%%%%%%%%%%%%%% Internal Review %%%%%%%%%%%%%%%%
% Paper Summary
% (I must preface this summary by noting that this topic is well outside of my wheelhouse, so apologies if I misrepresent some of the technical details.)

% An audio codec is a tool for compressing (encoding) and decompressing (decoding) digital audio, which is often used in digital speech processing systems (such as voice assistants or audio conferencing) to reduce network bandwidth requirements. Historically, audio codecs have been designed using psychoacoustics to determine which parts of the audio spectrum are (not) needed to preserve the perceived characteristics of music/speech. Recent work has eschewed this manual approach in favor of automated, learned approaches, such as the vector-quantized variational auto encoder (VQ-VAE). VQ-VAE has been demonstrated to work well for speech, but this paper argues that it might not be robust to speech "in the wild," where noise is an issue. The paper therefore argues for "enhancing the codec," which removes unwanted noise from the signal. Experiments with public and crowdsourced data demonstrate that the enhancing approach outperforms the baselines.

% Risks
% TL;DR: I don't see any obvious risk in publishing this paper, but I would like the authors to comment on one aspect (below).

% [Data] The utterances used for training/evaluation come from public datasets (LibriSpeech, AudioSet and VCTK), so no Amazon customer data is used.

% One thing I want to flag is that the mean opinion score (MOS) is computed using MTurk crowdsourcing. Since Amazon has paid to collect this data, it's unclear whether Amazon plans on releasing the data when the paper is published. If the conference wants all code/data sources to be released for reproducibility, this may become an issue.

% --- We are not obligated to release the MOS data nor do we plan to .

% [Tech] The ML model leans heavily on prior work (VQ-VAE), so I don't think any "secret sauce" is getting leaked. However, the paper gives a fairly detailed description of the model architecture. I'm not familiar with ICASSP papers, so not sure whether this level of detail is expected/common. What worries me is whether the paper describes the specific architecture of a production system, as this seems like "secret sauce" to me. I would like the authors to comment on this.

% --- The architectures evaluated and displayed in this is based off of open-source work and is not in production. The only production level system discussed here is the enhancement baseline whose architecture, training procedures, and dataset are not discussed here. 

% [Amazon Products] The paper does not mention any Amazon products or services other than Mechanical Turk, which is only used for evaluation. Thus, there is no risk of divulging forthcoming products.

% Feedback for Authors
% Again, I must preface these comments by saying that this is outside my area of expertise and experience, so I don't know what all is expected of ICASSP papers. All I can give is high-level feedback on the paper's presentation.

% In Section 2, I found it difficult to distinguish what was new (proposed) from what was existing work. Where does VQ-VAE stop and "compressor-enhancer" begin? I recommend clarifying this distinction.

% --- The distinction is based on the new joint compress-enhance task that we define and the corresponding training setup. The architectural advancement is minimal beyond how the encoder/decoder are interpreted.

% The terms "compressor-enhancer" and "enhancer-compressor" are both used. I recommend sticking with one term and being consistent about its usage.

% --- Thank you for catching this inconsistency. We have edited these to only be compressor-enhancer.

% I recommend moving Section 2.2 (Datasets) into Section 3. The proposed technique isn't really tied to any specific training/evaluation datasets; that's just what you used in these experiments.

% -- We appreciate this feedback and have moved the datasets section.
conclusion.tex
\section{Conclusion}
\label{sec:conclusion}
In this work we presented a model that does joint compression and enhancement of a noisy speech signal using a VQ-VAE with a convolutional encoder and a WaveRNN decoder. 
Through a set of mean opinion score based experiments, we found that joint compression and enhancement performs better in the presence of noise, including in low SNR scenarios, than stand-alone compression; and also outperforms a sequential combination of speech enhancement and a compression only neural codec. We also found that enhancement training improves codec performance on clean speech signals. 

% Future work on this topic could explore interactions with audio codecs of other audio processing tasks. 

% As shown, joint compression and enhancement is more robust to unseen noises, low SNR scenarios, and even clean speech than just compression. It also outperformed a sequential combination of speech enhancement and the compression only codec. 

%We combine prior research on neural audio codecs and on neural speech enhancement systems to produce a neural codec which simultaneously compresses and enhance a signal. 
%We call this process ``enhancing into the codec". 

% TODO Our results show that a compressor-enhancer model running at XX bps can XX.

% This line of research has many possible extensions and could be improved by more intricate mergers of compression and enhancement technologies. With this work we hope to show some initial promise for the task of joint compression-enhancement and inspire future work in the area.

% TODO cite: percepnet, Lim Chinen Jan paper, VQ-VAE codec paper, JM's 1.6kbps paper with Jan. 


% \begin{itemize}
%     \item We presented the task of joint enhancement and compression
%     \item More robust then just compression
%     \item Much smaller than sequential enhancement and compression
%     \item Exciting new set of problems
% \end{itemize}
exp_results.tex
\section{Experiments and Results}
\label{sec:exp_res}
We train the Codec Only and Enhancing Codec models at two different kb/s and for the highest kb/s Enhancing Codec we also experiment with modifying the training setup to use higher SNR mixtures.

%two, or three speech codebooks of size $512$. We do not vary the size of the speaker codebook. This corresponds to running models with bit-rates of approximately $.9$kb/s, and $1.35$kb/s.

\subsection{Model Details}
The encoder first computes an $80$ bin log-mel representation with a hop-size of $10$ms and a window size of $250$ms on $16$kHz audio. These are passed to the speech encoder which has five convolutional layers each with $768$ filters. The first, second, fourth and fifth layers use a stride of $1$ and a kernel of size $3$. The third layer downsamples by using a stride of $2$ and a kernel of size $4$. The speaker encoder has an identical architecture but uses $64$ filters and omits the fifth layer. These are both passed to separate VQ bottlenecks which apply a linear layer with output size $64$ before quantizing. The speaker encoding for an entire input file is quantized using a single code from a $9$-bit codebook, while the speech encoding is quantized using two $9$-bit codebooks for the $.9\,\text{kb/s}$ model and three $9$-bit codebooks for the $1.35\,\text{kb/s}$  model. When several codebooks are used, each codebook uses its own linear layer and the resulting output quantizations are stacked. In our current implementation these steps are non-causal, but can easily be made causal or with custom look-ahead by using causal convolutions and adapting speaker-encodings over time. The WaveRNN model's first GRU which produces the coarse representation has $192$ hidden nodes, and its fine-representation GRU has $896$ hidden nodes. 

All VQ codebooks are trained using the exponential moving average technique from \cite{van2017neural}. We train the models with a batch size of 80 per GPU and a sample length of 1 second using the Adam optimizer on 8 NVidia V100 GPUs for 3 days.

% \subsection{Baselines}
% For comparison we construct two alternative pipelines and run a third. The first is identical in architecture to the proposed model but is only trained on compressing and reconstructing the clean speech signal $\mathbf{x} = \mathbf{s}$. Our second comparison reuses the aforementioned compression model but also uses the pre-trained speech RNNoise enhancer \cite{valin2018hybrid} as a preprocessing step, feeding the enhanced signal to the compressor. We will refer to the first setup as ``Codec Only" and the second setup as ``Enhancement, then codec". Our third baseline is a pre-trained LPCNet vocoder.

% --- NEW LOCATION FOR DATASETS ---
\subsection{Datasets}

\subsubsection{Training}
\label{sec:data_train}
To generate a training mixture, we retrieve clean speech data from the LibriSpeech dataset in \cite{panayotov2015librispeech}, and noise data from AudioSet \cite{gemmeke2017audio}. When selecting noise clips from AudioSet, we avoid any clips with speech related tags. To increase the prevalence of challenging noise we sample noise clips with non-stationary noise more frequently. The noisy mixtures are created with a random SNR between $-5$ and $25$ dB. Finally, all room impulse responses are synthetically generated using the image-source method. For additional details consult \cite{isik2020poconet}.

\subsubsection{Evaluation}
To evaluate, we use the test mixtures from the VCTK dataset \cite{valentini2017noisy}. It contains mixtures with SNRs of $2.5$dB, $7.5$dB, $12.5$dB and $17.5$dB across a variety of speakers and noises. When evaluating our models on clean speech we use the clean speech samples from the VCTK test set.



\subsection{Subjective Quality Evaluation}
Since the compression based models in this paper resynthesize waveforms, their performance is not aptly measured by standard numerical metrics; we therefore measure model performance using a Mean Opinion Score (MOS) from a crowd-sourced study on Mechanical Turk that uses the P808 evaluation method \cite{P808}. 

\subsection{Results}
We compare the compressor-enhancer model with a compressor-only counterpart of identical architecture and size at bandwidths of $1.35\,\text{kb/s}$ and $0.9\,\text{kb/s}$ . The compressor only model is trained on the clean speech setup described in section~\ref{sec:data_train}. We also experiment with the combination of a speech enhancement model (RNNoise) \cite{valin2018hybrid} and the compression only model. We refer to these baselines as "Codec Only", and "Enhancement, then codec" respectively. As a final point of comparison we also evaluate a pre-trained LPCNet vocoder \cite{valin2019lpcnet}.

\subsubsection{MOS vs Bandwidth in Noisy Speech}
% We compared MOS scores on the VCTK test set for the following models across different bandwidths: a pure compression only model, which is a lower-compute version of \cite{garbacea2019low}, our enhancement-trained compression model, and a combination of speech enhancement model (RNNoise) \cite{valin2018hybrid} and the compression only model above. 
Fig.~\ref{fig:mos_v_bps} displays the MOS scores of the models at different bandwidths. Within the examined bandwidths, our enhancement-trained compression model~(blue diamonds and circle) scores about $.6$ MOS above its non-enhancing counterpart~(black square). Interestingly, sequential speech enhancement and compression scores worse than compression alone. We suspect this stems from the compression only model being susceptible to out of distribution errors. Finally, the LPCNet model serves as another comparison point for a compression only neural codec. The model denoted as "Enhancing Codec LN" is identical to "Enhancing Codec" in architecture but was trained on lower noise mixture with SNR ranging from $5$dB to $25$dB.

\begin{figure}
    \centering
    \includegraphics[scale=0.55]{figures/mos_vs_kbps.pdf}
    \vspace{-0.1in}
    \caption{MOS scores on the VCTK test set compared across models running at a range of kb/s. Our proposed joint compressor-enhancer models outperform the compression only baselines as well as the sequential enhancement-compression baselines. The LN suffix denotes a model trained on mixtures with lower noise content. The results are statistically significant with 95\% confidence intervals of approximately $.03$.}
    \label{fig:mos_v_bps}
\end{figure}


\subsubsection{Comparison of MOS across SNRs}
To see under what acoustic conditions enhancement-trained compression has the largest effect, we split MOS scores by SNR, and compare the $1.35\,\text{kb/s}$ versions of the three models discussed above. We display these results in Table \ref{Tab:diff_snr}. The MOS scores show that enhancement-trained compression compares favorably to both baselines across all SNRs. In the $2.5$dB scene the enhancing codec has a $.65$ to $.77$ lead over the baselines. This margin is reduced to $.38$ in the higher SNR scenes.

\begin{table}
    \begin{center}
            \vspace{-0.005in}
            \begin{tabular}{||c |c c c c||} 
            \hline
            Model \textbackslash SNR (dB) &  17.5 &  12.5 &  7.5 &  2.5 \\
            \hline\hline
             \thead{Codec Only}&  2.89&  2.79&  2.61&  2.34\\ 
            \hline
             \thead{Enhancing Codec LN}&  \textbf{3.27}&  \textbf{3.28}&  \textbf{3.22}&  \textbf{3.11}\\ 
            \hline
             \thead{Enhancement, then Codec}&  2.77&  2.64&  2.53&  2.46\\ 
            \hline
            \end{tabular}
    \end{center}
    \vspace{-0.1in}
    \caption{MOS comparison across SNRs on VCTK at $1.35\,\text{kb/s}$. The largest performance difference is in low SNR scenes.}
    \vspace{-0.1in}
    \label{Tab:diff_snr}
\end{table}



\subsubsection{Comparison of MOS on Clean Speech}
Observing the results at high-SNR, we also evaluated the MOS performance of enhancement-trained compression on clean speech. We compare the $1.35\,\text{kb/s}$ compression only model against our $1.35\,\text{kb/s}$ enhancement-trained compression model and display the results in Table \ref{Tab:clean}. We chose to omit the sequential enhancement then compression model since the speech is already clean. The enhancement-trained model outperforms the compression only model, leading us to suspect that training with noise helps the model learn a more robust bottleneck, and thus generalize better.

\begin{table}
    \begin{center}
            \begin{tabular}{||c | c||} 
            \hline
            Model&  Clean Speech MOS\\
            \hline\hline
             \thead{Codec Only}&  2.95\\ 
            \hline
             \thead{Enhancing Codec LN}&  \textbf{3.26}\\ 
            \hline
            \end{tabular}
    \end{center}
    \vspace{-0.1in}
    \caption{MOS comparison on VCTK clean speech at $1.35\,\text{kb/s}$. 95\% confidence intervals of $\approx .03$}
    \label{Tab:clean}
\end{table}




\subsection{Comparison with a two-stage approach}

% As shown, joint compression and enhancement is more robust to unseen noises, low SNR scenarios, and even clean speech than just compression. It also outperformed a sequential combination of speech enhancement and the compression only codec. 
We also attempted a joint compression and enhancement approach where we trained, first, a compression only model and then an encoder-only enhancement model trained to, given noisy speech, output the discretized latent representation of clean speech. With the goal being improvements to out-of-domain errors for speech enhancement models, as a two-stage approach would mean a fewer number of parameters needing to be trained on enhancement. However, we found that these models do not perform well, possibly because the clean-trained autoregressive decoder is too sensitive to out-of-domain inputs from the enhancer-encoder, making the simultaneous training of the decoder a key component of joint compression and enhancement.

% We hypothesized that by enhancing into the latent space you could circumvent artifacts common in mask-based enhancement model. We attempted to test this. However, we found this task difficult and sensitive to out of distribution errors leading to poor performance. We concluded that learning to compress jointly with learning to enhance was critical. This is corroborated by our experiments which showed that enhancement-trained compression outperforms sequential enhancement, then compression. While the compression performance of joint compression-enhancement is good, we think the enhancement performance could be further improved by incorporating speech enhancement advances such as waveform level models and specialized layer architectures. 



% \subsubsection{Independence of Encodings and Noise}
% In this final experiment we hypothesize that the encoder is performing enhancement into the quantized space and that a speech utterance will be encoded to the same codebooks even when mixed with noise at a different SNR.. To verify this we first encode a clean speech utterance and track the selected codebooks. Then, we mix the original utterance with noise at some SNR and check how many of the same codebooks were selected. We perform this analysis across a variety of SNRs and compare the enhancing and non-enhancing model. 
% % ANALYSIS BASED ON RESULTS HERE


% \begin{table}
%     \begin{center}
%             \begin{tabular}{||c |c c c c||} 
%             \hline
%             Model&  XdB&  XdB&  XdB&  XdB\\
%             \hline\hline
%              \thead{Non-Enhancing \\ Codec}&  0\%&  0\%&  0\%&  0\% \\ 
%             \hline
%              \thead{Enhancing \\ Codec}&  0\%&  0\%&  0\%&  0\% \\ 
%             \hline
%             \end{tabular}
%     \end{center}
%     \caption{Codebook proportion kept the same across SNRs.}
%     \label{Tab:cb_same}
% \end{table}

% \begin{itemize}
%     \item Joint enhance-compress is smaller and faster
%     \item Decoder is a generative model and acts as a speech prior
%     \item Decoder synthesizes unseen speakers well
%     \item Two stage learning is hard
%     \item Other setups to explore --  
% \end{itemize}



intro.tex
\section{Introduction}
\label{sec:intro}

Audio codecs compress speech signals by eliminating redundant and unnecessary information, with their design often leveraging extensive domain expertise to keep compression rates high, while keeping artifacts at a minimum. The most popular codecs, like the Opus codec in wideband mode, can produce high-quality speech compression at around $9$ kb/s \cite{valin2012definition}. Recently, there have been successful efforts in building learned codecs; starting with replacing the decoders with learned decoders for fixed encoders, which can operate as low as $2.4$ kb/s to $1.6$ kb/s \cite{kleijn2018wavenet, valin2019real}. These learned decoders leverage advances in speech synthesizing generative models such as WaveNet, WaveRNN, and LPCNet \cite{oord2016wavenet, lorenzo2018towards, valin2019lpcnet}. More recently, in \cite{garbacea2019low, zhen2020psychoacoustic}, the encoder and decoder were both learned in a joint fashion, by using quantized bottlenecks based on Vector-Quantized Variational Auto-Encoders~(VQ-VAE) \cite{van2017neural}, and soft-to-hard quantizers. The VQVAE based model improved at $1.6$ kb/s, on hand-designed encoders at $2.4$ kb/s.

% Previously
%More recently, in \cite{garbacea2019low}, the encoder and decoder were both learned in a joint fashion, by using quantized bottlenecks, based on Vector-Quantized Variational Auto-Encoders~(VQ-VAE) \cite{van2017neural}, improving, at $1.6$ kb/s, on hand-designed encoders at $2.4$ kb/s.

VQ-VAE models are auto-encoders where latent vectors are quantized using a learned vector quantization scheme. These discrete representations have been shown to have a good inductive bias for speech and perform well on unsupervised acoustic unit discovery tasks \cite{chorowski2019unsupervised, chen2020unsupervised}. VQ-VAE models are apt for low-bitrate compression, as an input can be represented by a sequence of discrete codebook vector indices, while the location of codebook vectors can be hard-coded.

Fully learned codecs like \cite{garbacea2019low, zhen2020psychoacoustic} open new avenues for learning based compression and demonstrate strong results for compressing clean speech. However, they are tightly coupled with speech content and do not perform well under the noisy conditions a codec might encounter in the wild. In this work, we focus on making them robust to speech corrupted with noise. In \cite{lim2020robust}, this issue was addressed by training a noise-robust feature extractor based on the Siamese learning paradigm, and then training a WaveNet model conditioned on those features. Motivated by the performance of modern neural speech enhancers in removing unwanted noise and reverberation from audio signals \cite{xia2013speech, xu2014regression, weninger2015speech, isik2020poconet, valin2018hybrid}, we combine the learning based compression and learning based enhancement paradigms. We call the resulting paradigm "enhancing into the codec". The proposed model is based on VQ-VAE with a WaveRNN decoder, and, trained end-to-end as a speech enhancer, can simultaneously compress and enhance noisy speech signals, independent of speaker identity.
% REVISIT
We refer to such a model as a compressor-enhancer: a model which jointly compresses and enhances speech.

%Our results spea
% The main contributions of this work are as follows:

% \begin{enumerate}
%     \item Introduction of the joint compression-enhancement paradigm
%     \item A robust technique for training neural audio codecs
%     \item An interpretation of neural audio codecs as a codec based enhancement step and a noise independent sysnthesis step
% \end{enumerate}

% Since the compression based models in this paper resynthesize waveforms from scratch, traditional reference-comparison based numerical metrics are not suitable for. 

We measure the performance of our models using Mean Opinion Scores (MOS) from a crowd-sourced study on Amazon Mechanical Turk. Across a range of compression rates and noise levels, we compare our model to a non-enhancing learned compressor both with and without additional enhancement preprocessing, as well as the LPCNet neural codec \cite{valin2019lpcnet}. We find that the proposed model performs well in noisy scenarios, compared to both non-enhancing codecs, as well as the composition, with comparable total compute cost, of a separate speech enhancer and a non-enhancing codec. We also find that the compressor-enhancer VQ-VAE performs significantly better than a clean-speech-trained VQ-VAE codec on clean speech inputs.

% \begin{itemize}
%     \item Codecs are hand tuned 
%     \item Neural Codecs are possible thanks to good synthesis models and VQ-VAE
%     \item Neural Codecs can compress clean speech really well
%     \item Neural Codecs are not robust to noise
%     \item Some SE background and citations
%     \item New paradigm: enhance into the compressed space
%     \item We run experiments to compare against a non-enhancing neural codec and analyze some codec properties
% \end{itemize}
method.tex

\section{Method}
\label{sec:method}

\begin{figure*}
    \centering
    \includegraphics[scale=0.75]{figures/block_diagram.pdf}
    \vspace{-3mm}
    \caption{Block Diagram of the compressor-enhancer. The speech and speaker encoders are made up of several convolutional layers with batch normalization and ReLU. The VQ bottleneck has separate quantizers for the speech and speaker encodings. In our experiments the speech quantizer is made up of several codebooks of size 512 and the speaker quantizer is made up of a single codebook of size 512. The decoder is a WaveRNN based model and uses the quantized speech and speaker information first to reconstruct a coarse and then reconstruct a fine waveform. The waveform output is mu-law quantized.}
    \vspace{-2mm}
    \label{fig:model}
\end{figure*}

The task of speech enhancement is to recover a clean speech signal $\mathbf{s}$ from a noisy and possibly reverberant mixture $\mathbf{x} = \mathbf{s} * \mathbf{h} + \mathbf{n}$ where $\mathbf{n}$ represents some additive noise signal and $\mathbf{s} * \mathbf{h}$ represents the convolution of a room impulse response $\mathbf{h}$ with the speech signal. The goal of speech compression is to reconstruct a speech signal  $\mathbf{s}$ after encoding it to a smaller representation $\mathbf{\hat{E}}$ and decoding it to the reconstruction $\mathbf{\hat{s}}$. A neural compression model is composed of an encoder network $\mathcal{N}_e$, a coding step $\mathcal{C}$ and decoder network $\mathcal{N}_d$ which balance a trade-off between reconstruction fidelity and the size of $\mathbf{\hat{E}}$.

Thus, we define the joint compression-enhancement task. In the joint compression-enhancement task, the model receives a noisy input $\mathbf{x}$, which it encodes to a smaller representation $\mathbf{\hat{E}}$ and then decodes to an estimated clean and decompressed speech signal $\mathbf{\hat{s}}$. The full procedure is therefore

\vspace{-3mm}
\begin{equation}
    \mathbf{\hat{s}} = \mathcal{N}_d( \overbrace{\mathcal{C}(\underbrace{\mathcal{N}_e(\mathbf{x})}_{\mathbf{E}})}^{\mathbf{\hat{E}}}).
\end{equation}
\vspace{-4mm}

\subsection{Model}
The proposed autoencoder model~(Fig. \ref{fig:model}) is comprised of a convolutional encoder, a VQ-VAE bottleneck, and a recurrent decoder. These correspond to $\mathcal{N}_e$, $\mathcal{C}$, and $\mathcal{N}_d$ respectively. The model takes as input a 16 kHz noisy speech signal which is processed by the encoder and quantized by the bottleneck. The decoder autoregressively reconstructs the original 16 kHz waveform using the quantized speech and speaker encodings. As such, the encoder is encouraged to produce a compressed representation that gives the most information for the decoder to conditionally model the clean speech signal.  
% While, in principle, this model could be run in real time, in our experiments, the encoder generates a single speaker encoding per input, making it non-causal.



\subsubsection{Encoder}
The encoder first computes a log-Mel representation and then applies a series of 1D convolutional layers, treating the mel bins as features. Each convolutional layer is followed by a batch normalization and then a ReLU non-linearity. The stride of the log-Mel representation and convolutional layers are selected to produce encodings at a rate of 50 Hz. The encoder also estimates one additional ``speaker embedding'' vector by performing a simple average across time over the output of a separate set of encoding layers. The output of theses steps is called $\mathbf{E}$. 

\subsubsection{Vector-Quantized Bottleneck}
The vector quantized bottleneck quantizes the outputs of the encoding layer using a set of codebooks. Where a separate codebook, constant over the entire input, is used to quantize the speaker embedding. We represent a codebook containing $K$ codes by $\mathbf{C} = \{\mathbf{c}_1, \cdots \mathbf{c}_K\}$. In the forward pass, the encoder outputs $\mathbf{E} = \{\mathbf{e}_1, \cdots \mathbf{e}_K\}$ are quantized by replacing each $\mathbf{e}_i$ with the closest $\mathbf{c}_j$ to get the quantized encoding $\mathbf{\hat{e}}_i$, where $j = \argmin_k ||\mathbf{e}_i - \mathbf{c}_k||_2^2$. Due to the non-differntiability of the $\argmin$ operation, VQ-VAE uses an additional two loss terms. The terms encourage each encoding $\mathbf{e}_i$ to be close to the selected $\mathbf{c}_j$ and for the each code $\mathbf{c}_j$ to minimize the quantization error incurred by any encodings that selected it. These are summarized below using the stopgradient operator $\operatorname{sg}$ which is identity at the forward pass but stops gradients in the reverse pass. In practice, we optimize the second term using an exponential moving average k-means. For additional details, see \cite{van2017neural}.

\begin{equation}
    \mathcal{L}_{\text{vq}} = \lambda ||\operatorname{sg}[\mathbf{\hat{E}}] - \mathbf{E}] ||_2^2 + ||\operatorname{sg}[\mathbf{E}] - \mathbf{\hat{E}}] ||_2^2.
\end{equation}

\subsubsection{Autoregressive Decoder}
We use an RNN based model to synthesize raw 16 kHz audio. The model, which is described in \cite{lorenzo2018towards}, contains two Gated Recurrent Units~(GRU) and two dense layers. We first concatenate the quantized speaker embedding to the quantized encoding and pass the resulting tensor through the first GRU. Then, we up-sample the GRU output to match the desired output length~(in raw audio samples) and pass the upsampled tensor through the second GRU, and two final dense layers. We apply softmax to the final dense layer and train the model to predict a distribution over 8-bit mu-law quantized values.

\subsubsection{Loss}
The final forward pass procedure is composed of passing raw noisy audio $\mathbf{x}$ to the encoder, quantizing the resulting encodings and speaker embedding, and running an autoregressive model to produce an estimated clean waveform $\mathbf{\hat{s}}$. The full loss function shown below is composed of $\mathcal{L}_{\text{vq}}$ and a cross-entropy term $\mathcal{L}_{\text{ce}}(\mathbf{s}, \mathbf{\hat{s}})$, which measures the KL-divergence between the predicted distribution and the one-hot value of the mu-law quantized clean speech $\mathbf{s}$.

\begin{equation}
    \mathcal{L} = \mathcal{L}_{\text{vq}} + \mathcal{L}_{\text{ce}}(\mathbf{s}, \mathbf{\hat{s}}).
\end{equation}


% MOVED TO THE NEXT SECTION
% \subsection{Datasets}

% \subsubsection{Training}
% \label{sec:data_train}
% To generate a training mixture, we retrieve clean speech data from the LibriSpeech dataset in \cite{panayotov2015librispeech}, and noise data from AudioSet \cite{gemmeke2017audio}. When selecting noise clips from AudioSet, we avoid any clips with speech related tags. To increase the prevalence of challenging noise we sample noise clips with non-stationary noise more frequently. The noisy mixtures are created with a random SNR between $-5$ and $25$ dB. Finally, all room impulse responses are synthetically generated using the image-source method. For additional details consult \cite{isik2020poconet}.

% \subsubsection{Evaluation}
% To evaluate, we use the test mixtures from the VCTK dataset \cite{valentini2017noisy}. It contains mixtures with SNRs of $2.5$dB, $7.5$dB, $12.5$dB and $17.5$dB across a variety of speakers and noises. When evaluating our models on clean speech we use the clean speech samples from the VCTK test set.



main.bbl
\begin{thebibliography}{10}

\bibitem{valin2012definition}
Jean-Marc Valin, Koen Vos, and Timothy Terriberry,
\newblock ``Definition of the opus audio codec,''
\newblock {\em IETF, September}, 2012.

\bibitem{kleijn2018wavenet}
W~Bastiaan Kleijn, Felicia~SC Lim, Alejandro Luebs, Jan Skoglund, Florian
  Stimberg, Quan Wang, and Thomas~C Walters,
\newblock ``Wavenet based low rate speech coding,''
\newblock in {\em 2018 IEEE International Conference on Acoustics, Speech and
  Signal Processing (ICASSP)}. IEEE, 2018, pp. 676--680.

\bibitem{valin2019real}
Jean-Marc Valin and Jan Skoglund,
\newblock ``A real-time wideband neural vocoder at 1.6 kb/s using lpcnet,''
\newblock {\em Proc. Interspeech 2019}, pp. 3406--3410, 2019.

\bibitem{oord2016wavenet}
Aaron van~den Oord, Sander Dieleman, Heiga Zen, Karen Simonyan, Oriol Vinyals,
  Alex Graves, Nal Kalchbrenner, Andrew Senior, and Koray Kavukcuoglu,
\newblock ``Wavenet: A generative model for raw audio,''
\newblock {\em arXiv preprint arXiv:1609.03499}, 2016.

\bibitem{lorenzo2018towards}
Jaime Lorenzo-Trueba, Thomas Drugman, Javier Latorre, Thomas Merritt, Bartosz
  Putrycz, Roberto Barra-Chicote, Alexis Moinet, and Vatsal Aggarwal,
\newblock ``Towards achieving robust universal neural vocoding,''
\newblock {\em arXiv preprint arXiv:1811.06292}, 2018.

\bibitem{valin2019lpcnet}
Jean-Marc Valin and Jan Skoglund,
\newblock ``Lpcnet: Improving neural speech synthesis through linear
  prediction,''
\newblock in {\em ICASSP 2019-2019 IEEE International Conference on Acoustics,
  Speech and Signal Processing (ICASSP)}. IEEE, 2019, pp. 5891--5895.

\bibitem{garbacea2019low}
Cristina G{\^a}rbacea, A{\"a}ron van~den Oord, Yazhe Li, Felicia~SC Lim,
  Alejandro Luebs, Oriol Vinyals, and Thomas~C Walters,
\newblock ``Low bit-rate speech coding with vq-vae and a wavenet decoder,''
\newblock in {\em ICASSP 2019-2019 IEEE International Conference on Acoustics,
  Speech and Signal Processing (ICASSP)}. IEEE, 2019, pp. 735--739.

\bibitem{zhen2020psychoacoustic}
Kai Zhen, Mi~Suk Lee, Jongmo Sung, Seungkwon Beack, and Minje Kim,
\newblock ``Psychoacoustic calibration of loss functions for efficient
  end-to-end neural audio coding,''
\newblock {\em IEEE Signal Processing Letters}, vol. 27, pp. 2159--2163, 2020.

\bibitem{van2017neural}
Aaron Van Den~Oord, Oriol Vinyals, et~al.,
\newblock ``Neural discrete representation learning,''
\newblock in {\em Advances in Neural Information Processing Systems}, 2017, pp.
  6306--6315.

\bibitem{chorowski2019unsupervised}
Jan Chorowski, Ron~J Weiss, Samy Bengio, and A{\"a}ron van~den Oord,
\newblock ``Unsupervised speech representation learning using wavenet
  autoencoders,''
\newblock {\em IEEE/ACM transactions on audio, speech, and language
  processing}, vol. 27, no. 12, pp. 2041--2053, 2019.

\bibitem{chen2020unsupervised}
Mingjie Chen and Thomas Hain,
\newblock ``Unsupervised acoustic unit representation learning for voice
  conversion using wavenet auto-encoders,''
\newblock {\em arXiv preprint arXiv:2008.06892}, 2020.

\bibitem{lim2020robust}
Felicia~SC Lim, W~Bastiaan Kleijn, Michael Chinen, and Jan Skoglund,
\newblock ``Robust low rate speech coding based on cloned networks and
  wavenet,''
\newblock in {\em ICASSP 2020-2020 IEEE International Conference on Acoustics,
  Speech and Signal Processing (ICASSP)}. IEEE, 2020, pp. 6769--6773.

\bibitem{xia2013speech}
Bingyin Xia and Changchun Bao,
\newblock ``Speech enhancement with weighted denoising auto-encoder.,''
\newblock in {\em Interspeech}, 2013, pp. 3444--3448.

\bibitem{xu2014regression}
Yong Xu, Jun Du, Li-Rong Dai, and Chin-Hui Lee,
\newblock ``A regression approach to speech enhancement based on deep neural
  networks,''
\newblock {\em IEEE/ACM Transactions on Audio, Speech, and Language
  Processing}, vol. 23, no. 1, pp. 7--19, 2014.

\bibitem{weninger2015speech}
Felix Weninger, Hakan Erdogan, Shinji Watanabe, Emmanuel Vincent, Jonathan
  Le~Roux, John~R Hershey, and Bj{\"o}rn Schuller,
\newblock ``Speech enhancement with lstm recurrent neural networks and its
  application to noise-robust asr,''
\newblock in {\em International Conference on Latent Variable Analysis and
  Signal Separation}. Springer, 2015, pp. 91--99.

\bibitem{isik2020poconet}
Umut Isik, Ritwik Giri, Neerad Phansalkar, Jean-Marc Valin, Karim Helwani, and
  Arvindh Krishnaswamy,
\newblock ``Poconet: Better speech enhancement with frequency-positional
  embeddings, semi-supervised conversational data, and biased loss,''
\newblock {\em arXiv preprint arXiv:2008.04470}, 2020.

\bibitem{valin2018hybrid}
Jean-Marc Valin,
\newblock ``A hybrid dsp/deep learning approach to real-time full-band speech
  enhancement,''
\newblock in {\em 2018 IEEE 20th International Workshop on Multimedia Signal
  Processing (MMSP)}. IEEE, 2018, pp. 1--5.

\bibitem{panayotov2015librispeech}
Vassil Panayotov, Guoguo Chen, Daniel Povey, and Sanjeev Khudanpur,
\newblock ``Librispeech: an asr corpus based on public domain audio books,''
\newblock in {\em 2015 IEEE International Conference on Acoustics, Speech and
  Signal Processing (ICASSP)}. IEEE, 2015, pp. 5206--5210.

\bibitem{gemmeke2017audio}
Jort~F Gemmeke, Daniel~PW Ellis, Dylan Freedman, Aren Jansen, Wade Lawrence,
  R~Channing Moore, Manoj Plakal, and Marvin Ritter,
\newblock ``Audio set: An ontology and human-labeled dataset for audio
  events,''
\newblock in {\em 2017 IEEE International Conference on Acoustics, Speech and
  Signal Processing (ICASSP)}. IEEE, 2017, pp. 776--780.

\bibitem{valentini2017noisy}
Cassia Valentini-Botinhao et~al.,
\newblock ``{Noisy speech database for training speech enhancement algorithms
  and TTS models},''
\newblock {\em University of Edinburgh. School of Informatics. Centre for
  Speech Technology Research (CSTR)}, 2017.

\bibitem{P808}
ITU-T,
\newblock {\em Recommendation {P}.808: Subjective evaluation of speech quality
  with a crowdsourcing approach}, 2018.

\end{thebibliography}

strings.bib
@inproceedings{garbacea2019low,
  title={Low bit-rate speech coding with VQ-VAE and a WaveNet decoder},
  author={G{\^a}rbacea, Cristina and van den Oord, A{\"a}ron and Li, Yazhe and Lim, Felicia SC and Luebs, Alejandro and Vinyals, Oriol and Walters, Thomas C},
  booktitle={ICASSP 2019-2019 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)},
  pages={735--739},
  year={2019},
  organization={IEEE}
}

@inproceedings{van2017neural,
  title={Neural discrete representation learning},
  author={Van Den Oord, Aaron and Vinyals, Oriol and others},
  booktitle={Advances in Neural Information Processing Systems},
  pages={6306--6315},
  year={2017}
}

@article{kalchbrenner2018efficient,
  title={Efficient neural audio synthesis},
  author={Kalchbrenner, Nal and Elsen, Erich and Simonyan, Karen and Noury, Seb and Casagrande, Norman and Lockhart, Edward and Stimberg, Florian and Oord, Aaron van den and Dieleman, Sander and Kavukcuoglu, Koray},
  journal={arXiv preprint arXiv:1802.08435},
  year={2018}
}

@article{lorenzo2018towards,
  title={Towards achieving robust universal neural vocoding},
  author={Lorenzo-Trueba, Jaime and Drugman, Thomas and Latorre, Javier and Merritt, Thomas and Putrycz, Bartosz and Barra-Chicote, Roberto and Moinet, Alexis and Aggarwal, Vatsal},
  journal={arXiv preprint arXiv:1811.06292},
  year={2018}
}

@inproceedings{kleijn2018wavenet,
  title={WaveNet based low rate speech coding},
  author={Kleijn, W Bastiaan and Lim, Felicia SC and Luebs, Alejandro and Skoglund, Jan and Stimberg, Florian and Wang, Quan and Walters, Thomas C},
  booktitle={2018 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)},
  pages={676--680},
  year={2018},
  organization={IEEE}
}

@inproceedings{lim2020robust,
  title={Robust Low Rate Speech Coding Based on Cloned Networks and Wavenet},
  author={Lim, Felicia SC and Kleijn, W Bastiaan and Chinen, Michael and Skoglund, Jan},
  booktitle={ICASSP 2020-2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)},
  pages={6769--6773},
  year={2020},
  organization={IEEE}
}

@article{zhen2020psychoacoustic,
  title={Psychoacoustic Calibration of Loss Functions for Efficient End-to-End Neural Audio Coding},
  author={Zhen, Kai and Lee, Mi Suk and Sung, Jongmo and Beack, Seungkwon and Kim, Minje},
  journal={IEEE Signal Processing Letters},
  volume={27},
  pages={2159--2163},
  year={2020},
  publisher={IEEE}
}

@article{oord2016wavenet,
  title={Wavenet: A generative model for raw audio},
  author={Oord, Aaron van den and Dieleman, Sander and Zen, Heiga and Simonyan, Karen and Vinyals, Oriol and Graves, Alex and Kalchbrenner, Nal and Senior, Andrew and Kavukcuoglu, Koray},
  journal={arXiv preprint arXiv:1609.03499},
  year={2016}
}

@article{van2020vector,
  title={Vector-quantized neural networks for acoustic unit discovery in the ZeroSpeech 2020 challenge},
  author={van Niekerk, Benjamin and Nortje, Leanne and Kamper, Herman},
  journal={arXiv preprint arXiv:2005.09409},
  year={2020}
}

@inproceedings{xia2013speech,
  title={Speech enhancement with weighted denoising auto-encoder.},
  author={Xia, Bingyin and Bao, Changchun},
  booktitle={Interspeech},
  pages={3444--3448},
  year={2013}
}

@article{xu2014regression,
  title={A regression approach to speech enhancement based on deep neural networks},
  author={Xu, Yong and Du, Jun and Dai, Li-Rong and Lee, Chin-Hui},
  journal={IEEE/ACM Transactions on Audio, Speech, and Language Processing},
  volume={23},
  number={1},
  pages={7--19},
  year={2014},
  publisher={IEEE}
}

@inproceedings{weninger2015speech,
  title={Speech enhancement with LSTM recurrent neural networks and its application to noise-robust ASR},
  author={Weninger, Felix and Erdogan, Hakan and Watanabe, Shinji and Vincent, Emmanuel and Le Roux, Jonathan and Hershey, John R and Schuller, Bj{\"o}rn},
  booktitle={International Conference on Latent Variable Analysis and Signal Separation},
  pages={91--99},
  year={2015},
  organization={Springer}
}

@article{han2015learning,
  title={Learning spectral mapping for speech dereverberation and denoising},
  author={Han, Kun and Wang, Yuxuan and Wang, DeLiang and Woods, William S and Merks, Ivo and Zhang, Tao},
  journal={IEEE/ACM Transactions on Audio, Speech, and Language Processing},
  volume={23},
  number={6},
  pages={982--992},
  year={2015},
  publisher={IEEE}
}

@article{isik2020poconet,
  title={PoCoNet: Better Speech Enhancement with Frequency-Positional Embeddings, Semi-Supervised Conversational Data, and Biased Loss},
  author={Isik, Umut and Giri, Ritwik and Phansalkar, Neerad and Valin, Jean-Marc and Helwani, Karim and Krishnaswamy, Arvindh},
  journal={arXiv preprint arXiv:2008.04470},
  year={2020}
}

@article{chen2020unsupervised,
  title={Unsupervised acoustic unit representation learning for voice conversion using WaveNet auto-encoders},
  author={Chen, Mingjie and Hain, Thomas},
  journal={arXiv preprint arXiv:2008.06892},
  year={2020}
}

@article{chorowski2019unsupervised,
  title={Unsupervised speech representation learning using wavenet autoencoders},
  author={Chorowski, Jan and Weiss, Ron J and Bengio, Samy and van den Oord, A{\"a}ron},
  journal={IEEE/ACM transactions on audio, speech, and language processing},
  volume={27},
  number={12},
  pages={2041--2053},
  year={2019},
  publisher={IEEE}
}

@inproceedings{maiti2020speaker,
  title={Speaker independence of neural vocoders and their effect on parametric resynthesis speech enhancement},
  author={Maiti, Soumi and Mandel, Michael I},
  booktitle={ICASSP 2020-2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)},
  pages={206--210},
  year={2020},
  organization={IEEE}
}

@inproceedings{panayotov2015librispeech,
  title={Librispeech: an asr corpus based on public domain audio books},
  author={Panayotov, Vassil and Chen, Guoguo and Povey, Daniel and Khudanpur, Sanjeev},
  booktitle={2015 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)},
  pages={5206--5210},
  year={2015},
  organization={IEEE}
}

@article{valentini2017noisy,
  title={{Noisy speech database for training speech enhancement algorithms and TTS models}},
  author={Valentini-Botinhao, Cassia and others},
  journal={University of Edinburgh. School of Informatics. Centre for Speech Technology Research (CSTR)},
  year={2017}
}

@article{reddy2020interspeech,
  title={The interspeech 2020 deep noise suppression challenge: Datasets, subjective speech quality and testing framework},
  author={Reddy, Chandan KA and Beyrami, Ebrahim and Dubey, Harishchandra and Gopal, Vishak and Cheng, Roger and Cutler, Ross and Matusevych, Sergiy and Aichner, Robert and Aazami, Ashkan and Braun, Sebastian and others},
  journal={arXiv preprint arXiv:2001.08662},
  year={2020}
}

@article{valin2007speex,
  title={The speex codec manual version 1.2 beta 3},
  author={Valin, Jean-Marc},
  journal={Xiph. org Foundation},
  year={2007}
}

@article{valin2012definition,
  title={Definition of the Opus audio codec},
  author={Valin, Jean-Marc and Vos, Koen and Terriberry, Timothy},
  journal={IETF, September},
  year={2012}
}

@article{valin2019real,
  title={A Real-Time Wideband Neural Vocoder at 1.6 kb/s Using LPCNet},
  author={Valin, Jean-Marc and Skoglund, Jan},
  journal={Proc. Interspeech 2019},
  pages={3406--3410},
  year={2019}
}

@inproceedings{valin2019lpcnet,
  title={LPCNet: Improving neural speech synthesis through linear prediction},
  author={Valin, Jean-Marc and Skoglund, Jan},
  booktitle={ICASSP 2019-2019 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)},
  pages={5891--5895},
  year={2019},
  organization={IEEE}
}

@inproceedings{valin2018hybrid,
  title={A hybrid DSP/deep learning approach to real-time full-band speech enhancement},
  author={Valin, Jean-Marc},
  booktitle={2018 IEEE 20th International Workshop on Multimedia Signal Processing (MMSP)},
  pages={1--5},
  year={2018},
  organization={IEEE}
}

@inproceedings{gemmeke2017audio,
  title={Audio set: An ontology and human-labeled dataset for audio events},
  author={Gemmeke, Jort F and Ellis, Daniel PW and Freedman, Dylan and Jansen, Aren and Lawrence, Wade and Moore, R Channing and Plakal, Manoj and Ritter, Marvin},
  booktitle={2017 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)},
  pages={776--780},
  year={2017},
  organization={IEEE}
}

@article{naderi2020impact,
  title={Impact of the Number of Votes on the Reliability and Validity of Subjective Speech Quality Assessment in the Crowdsourcing Approach},
  author={Naderi, Babak and Hossfeld, Tobias and Hirth, Matthias and Metzger, Florian and M{\"o}ller, Sebastian and Jim{\'e}nez, Rafael Zequeira},
  journal={arXiv preprint arXiv:2003.11300},
  year={2020}
}

@MANUAL {P808,
        author = "ITU-T",
        title = "Recommendation {P}.808: Subjective evaluation of speech quality with a crowdsourcing approach",
        year = "2018",
}
