Neural Audio Codecを提案する論文を最近よく見る気がしますね。音楽生成系の研究と密接に関連するトピックなので、これまで提案されている手法を一通り整理してみようと思います。

（WavTokenizer論文によるコーデック性能分布図）

基本的な枠組み：ベクトル量子化とVQ-VAE

少なくともここで挙げるコーデックは、いずれもベクトル量子化（Vector Quantization）手法を用いて音声データをエンコードしています。DNNオートエンコーダーによって圧縮された潜在特徴ベクトル列を、更に離散値（＝整数値）に落とし込むことで、デコード品質を保ったまま効率よくビットレートを落とすことができます。

VQ-VAEの学習は、デコーダー出力とエンコーダー出力を近付ける再構成ロスと、量子化コードブックの学習を安定化させるためのコミットロスの最適化からなります。ベクトル量子化部の安定的な学習は比較的難しく、色んなテクニックが考案されています。

ベクトル量子化の改善テクニックは、vector-quantize-pytorchという神ライブラリにだいたい実装されています。
https://github.com/lucidrains/vector-quantize-pytorch

Transformerに基づく自然言語生成モデルが大きな成功を収めると、当初はデータサイズを圧縮するために考案された量子化コードも自然言語のように「言語モデル」が構築され始め、コーデックは音楽生成モデルにとっての「トークナイザ（Tokenizer）」として利用され始めます。NLPの文脈に従い、量子化された離散値も「トークン」と呼ばれるようになりました。

Jukebox：標準的なVQ-VAE

2020年、（まだOpenだった）OpenAIが発表した音楽生成モデルです。Jukebox自体はNeural Audio Codecではないのですが、音楽オーディオをVQ-VAEでトークンに変換し、そのトークン列の生成過程をモデリングするというアプローチを採用しました。VQ-VAEは元々画像生成の分野で提案されたものですが、音楽データに初めて適用したのはJukeboxだと思われます。

Jukeboxが提案したトークン化手法は標準的なVQ-VAEですが、生成モデルにおける使い方が少々独特です。下の図のように、異なる時間解像度のVQ-VAEを作り、生成モデルに利用しています。

Jukeboxは、時間解像度が最も低い（＝トークン数が一番少ない）トークン列の生成分布とともに、「低解像度トークンを高解像度トークンに変換する」条件分布を学習し、
低解像度トークン生成⇒高解像度化⇒高解像度化
という順番で推論を行うことで、高品質な音楽の生成を実現しました。

EnCodec：ベクトル量子化の階層化拡張

Meta AIが提案したニューラルコーデックです。主な構造はGoogleが先に考案したSoundStreamとほぼ全く同じですが、オープンソースでモデルチェックポイントも公開されていることから、今ではNeural Audio Codecのスタンダードというか、主なベースラインと見なされており、音楽生成モデルMusicGenをはじめ様々な生成モデルで使われています。

SoundStreamやEnCodecの主な特徴である「Residual Vector Quantization (RVQ)」は、トークン列を多層化することで、トークン列の情報量を効率的に向上させる手法です。具体的には、ベクトル量子化の量子化差分に再度ベクトル量子化を行い、それを繰り返すことでベクトルを複数層の量子化コードに変換します。

このような構造にすることで、Residual VQは低いビットレートで高い表現力を持てるだけでなく、層を重ねるごとに高精度化するような階層構造をトークン列に持たせることができます。なので、Jukeboxがやったような「低解像度トークンを生成→高解像度化」の生成アプローチを適用することも可能です。

SoundStreamとEnCodecはさらにHifiGANなどのボコーダー学習のノウハウをデコーダー学習に取り入れ、マルチ解像度のメルスペクトログラム再構成ロスを使ったり、異なる解像度のSTFTに対するDiscriminatorを用意して敵対的学習を行ったりして、音質向上を図りました。

（データのドメインを一層偏らせることにもなると思いますが…）

EnCodec独自の工夫としては、まずRVQの各階層が特定の周波数帯域に対応するようにしました。たとえば24kHzの信号を5層のRVQにエンコードする場合、それぞれの階層は1.5、3、6、12、24kHzまでの帯域に対応します。学習時は、ある帯域のデータで訓練するときは対応する階層まで使う、という感じで、ネットワークを学習させていきます。

また、軽量なTransformer「言語モデル」を用いたRVQトークンの再圧縮手法を提案しています。更に、敵対的学習を安定化させるための「loss balancer」提案し、Discriminatorから返ってくる勾配のスケールを制御することで、学習の安定性を大きく改善させました。

DAC：最新知見を借りた地道な改善

こちらも色んな所で利用されている、オープンソースのコーデックです。EnCodecをベースにしており、EnCodecが持つ問題を改善するため、色んな所から最新知見を持ってきて応用しました。論文ではそれぞれの効果が詳しく検証されています。

EnCodecのエンコーダー・デコーダーで採用されている、LeakyReluのような活性化関数は、周期性をもつ信号の生成に不向きなため、生成音の品質に影響を及ぼしていました。DACはBigVGANで採用されたSnakeという活性化関数を採用することで、音質を顕著に改善しています。

ベクトル量子化にありがちな問題として、コードブック内のごく一部のコードしか利用されない現象（codebook collapse）があります。SoundStreamやEnCodecでも、この現象を抑えるために、k-means法によるコードブック初期化、移動平均（EMA）でコードブック更新、低使用頻度コードのランダム初期化などのテクニックを導入していますが、DACは画像生成分野で提案された「factorized codes」と「L2-normalized codes」を採用することでコードブックの利用率を大きく向上させました。

「factorized codes」は、高い次元数（たとえば1024次元）を持つコードブックを、低い次元数（8か32次元）のベクトルからの線形投影で表す手法です。コードブックの次元数自体は高く保ったまま、潜在特徴とコードブックのマッチングは低次元空間で行うことで、崩壊を起こしにくくします。

「L2-normalized codes」は、潜在特徴とコードブックにL2正規化を行うことで、潜在空間内の距離尺度をユークリッド距離からコサイン距離に変換する手法です。これも学習の安定化に繋がることが示されています。

https://arxiv.org/abs/2110.04627

SoundStreamではQuantizer dropoutという、学習時は毎回ランダムに選ばれた一部のコードブックしか使用しないというテクニックを使っていましたが、DACではこれを改良し、Quantizer dropout自体も確率的に適用することで、音質の低下を抑えました。

敵対的学習に使うDiscriminatorも見直しています。EnCodecではSTFT振幅に対するDiscriminatorを使いましたが、DACは複素数STFTのDiscriminatorを使いました。

再構成ロスに使用されている、マルチ解像度のメルスペクトログラム再構成ロス・線形スペクトログラムの再構成ロスの細かい設定や重みも色々探索し、ベストな設定を見つけたそうです。

色々な改善を積み重ねた結果、実験ではEnCodecと同等の品質を更に低い圧縮率で実現しています。