wav2vec2.0を改造して音楽の特徴量を作成する

# はじめに
wav2vec 2.0は、音声データから特徴量を抽出するための自己教師あり学習モデルです。
自己教師あり学習は、ラベル付けされていないデータを利用してモデルを事前学習する手法で、ラベル付きデータを大量に用意できない場合に有効な手法です。
音声認識データセットの53,000時間のラベルなしデータで事前学習を行い、1時間のラベル付きデータでファインチューニングを行った結果、100時間のラベル付きデータで学習した従来の手法の性能を上回る結果を達成しています。

今回の目的は、このフレームワークを利用して、少量のデータで高い精度を出すための音楽特徴量を作成することです。
音楽のデータセットを作成するには専門家が必要であり、知識とかなりの労力が必要で、大量のデータを作成することが困難です。さらに音楽と一口に言っても、コードやビート、キーなど様々な要素があるため、多くの異なるタスクが存在します。
このようなタスクに対応するために汎用性の高い音楽の特徴量を作成したいと思います。

wav2vec2.0は53,000時間という驚異的なサイズのデータセットを使用していますが、個人でそのサイズのデータを学習させるのは困難なので、もう少し小規模なデータを事前学習に使用しています。さらにいくつかの変更点を加えており、主に計算量の削減と学習時間の削減を目的とした変更を行っています。

コードは[GitHub](https://github.com/anime-song/music-encoder)で公開しています。

# wav2vec2.0の構造
wav2vec 2.0は、特徴エンコーダとコンテキストネットワークの2つのコンポーネントから構成されています。
1. 音声データ$\mathcal{X}$を特徴エンコーダーで$\mathcal{Z}$に変換します。
2. $\mathcal{Z}$を量子化層で$\mathcal{Q}$に変換します。
3. $\mathcal{Z}$を一部マスクし、コンテキストネットワークに入力し、$\mathcal{C}$に変換します。
4. マスクをしたタイムステップである$\mathcal{C}_t$と$\mathcal{Q}_t$で対照損失を計算します。
![wav2vec2.0の参考画像](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1275467/7cbc9d75-c0ee-a979-dd52-aed7cd38c0c6.png)

### 特徴エンコーダー
特徴エンコーダは1次元の畳み込みレイヤーを利用し、生の音声データから音声表現を抽出します。
ストライド(5,2,2,2,2,2,2)とカーネルサイズ(10,3,3,3,3,2,2)を持ち、16kHzのデータを約320分の1の49Hzまでダウンサンプリングします。
特徴エンコーダーの目的は次元削減や有用な潜在表現を抽出することにあります。

### 量子化層
量子化層は、特徴エンコーダの出力の離散表現を学習します。
今回の変更では量子化層は安定性や学習速度に問題があるため線形層に置き換えています。

### コンテキストネットワーク
コンテキストネットワークはマスクされた特徴エンコーダーの出力を入力とし、マスクされたタイムステップの量子化層の出力を予測するように学習されます。(量子化層はマスキングされません)
マスキングは0.065の確率で開始位置を選択し、開始位置から10だけマスクされます。

### 対照損失 (Contrastive Loss)
コンテキストネットワークが量子化層の出力を予測するための学習に使用する損失関数は対照損失で、以下のように表されます。[^1]

```math
\mathcal{L}_m = -\log \frac{\exp(sim({\bf c}_t, {\bf q}_t) / \kappa)}{\sum_{\tilde{{\bf q}} \sim {\bf Q}_t} \exp(sim({\bf c}_t,\tilde{\bf q}) / \kappa)}
```
ここで、${\bf q}_t$と${\bf c}_t$はマスクされたタイムステップの量子化層とコンテキストネットワークの出力であり、$\kappa$は温度パラメータです。
分子部分はポジティブペアの類似度を、分母部分はネガティブペアの類似度とポジティブペアの類似度の合計を表しています。
$\tilde{\bf q}$は真の量子化表現(ポジティブサンプル)である${\bf q}_t$とK個のディストラクター（ネガティブサンプル）からなる集合からサンプリングされます。
Kは論文では100が指定されています。

具体例として、マスクされたタイムステップの${\bf c}_1$と${\bf q}_1$を考えます。
分子では${\bf c}_1$と${\bf q}_1$の類似度を計算し、分母では${\bf c}_1$と$\tilde{\bf q}$の類似度を計算し合計します。$\tilde{\bf q}$は${\bf q}_1$以外のマスクされたタイムステップからサンプリングされたネガティブサンプルと${\bf q}_1$です。

ここでネガティブペアはポジティブペアのペナルティのように動作します。仮にネガティブペアがなければ、すべて同じ出力にするだけで簡単にロスが最小化されてしまいます。ポジティブペアは量子化層とコンテキストネットワークの出力を近づけ、ネガティブペアはマスク部分と他のマスク部分の類似度を下げ、より情報を含んだ出力を促します。

実装例は以下のようになります。
```python
def contrastive_loss(self, mask, context, quantized):
    loss = []
    for batch_index in range(self.batch_size):
        # マスク部分のみを抜き出す
        mask_indices = tf.where(mask[batch_index])
        c_t = tf.gather_nd(context[batch_index], mask_indices)
        q_t = tf.gather_nd(quantized[batch_index], mask_indices)

        # ポジティブペアの計算
        pos_similarity = cosine_similarity(c_t, q_t)

        # ネガティブペアの計算
        neg_similarity = cosine_similarity_matmul(c_t, q_t)
        # ネガティブサンプルにはポジティブペアが含まれる可能性があるのでマスクする
        neg_similarity_mask = tf.linalg.diag(
            tf.ones(tf.shape(neg_similarity)[-1]))
        neg_similarity = tf.exp(
            neg_similarity / self.temperature) * (1 - neg_similarity_mask)

        # ランダムに並び替えてnum_negative_samplesだけ取得
        # この時ポジティブペアが選択されてしまった場合、
        # num_negative_samples-1のサンプル数になってしまいますが、大した影響はないと思います
        random_indices = tf.random.shuffle(
            tf.range(tf.shape(neg_similarity)[-1]))[:self.num_negative_samples]
        neg_similarity = tf.gather(
            neg_similarity, random_indices, axis=1)

        # ネガティブサンプルにポジティブペアを結合する
        neg_similarity = tf.concat(
                [neg_similarity, pos_similarity[:, tf.newaxis]], axis=-1)

        numerator = tf.exp(pos_similarity / self.temperature)
        denominator = tf.reduce_sum(masked_neg_similarity, axis=-1)
        loss_m = -tf.math.log(numerator / denominator)
        loss.append(tf.reduce_mean(loss_m))

    return tf.reduce_mean(loss)
```

# wav2vec2.0からの変更点
### 音響情報を捉えるための損失関数を追加
[AudioLM](https://arxiv.org/pdf/2209.03143.pdf)で言及されているようにwav2vec2.0のような学習方法で学習したモデルは音響情報をあまり含んでいません。音楽では音響情報は非常に重要で、特にコードやキーの予測には音響情報が必要不可欠です。
そこで[HuBERT](https://arxiv.org/pdf/2106.07447.pdf)を参考に対照損失とは別に音響情報を捉えるため損失関数を追加します。
具体的には入力を量子化し、コンテキストネットワークでマスクされた部分の量子化IDを予測するようにします。
HuBERTでは量子化はK-Meansで行われますが、今回はRVQ(Residual Vector Quantization)を量子化に使用します。

![HuBERTの参考画像](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1275467/415f9fec-4eb6-c842-7481-d62c455858dd.png)

### 入力データを生の音声データからスペクトログラムへ変更
wav2vec 2.0は元々生の音声データを入力として使用していますが、これは16kHzの1秒の音声データで16000の長さのデータとなり、メモリや計算量の観点から効率的ではないです。
そこで、入力データをスペクトログラムに変更しました。これに伴い、特徴エンコーダはストライド1のカーネルサイズ4に変更し、時系列のダウンサンプリングは行っていません。
1フレームあたりの秒数は論文と同じ0.02秒(20ms)に設定しています。
この変更による精度への影響は検証していませんが、同様にスペクトログラムを使用してwav2vec2.0を学習させている研究が存在します。[^4]

### 量子化層を線形層に変更
量子化層は勾配が非常に不安定で崩壊しやすい問題や、コードブックを均一に利用するためにDiversityLossを導入しているため、モデルが複雑になっている問題があります。
量子化層を線形層に置き換えることで、これらの問題を緩和します。

### コンテキストネットワークをDilated transformerに変更
wav2vec 2.0のコンテキストネットワークは、元々はTransformerを利用していますが、Transformerは計算量が入力の長さに対して2次的に増加するという欠点があります。
この問題を解決するために、[Beat Transformer](https://arxiv.org/pdf/2209.07140.pdf)で提案されたDilated transformerを利用しました。
Dilated transformerは計算量を抑えつつ、長期依存関係を効率的に捉えることができるようです。


# 実装
### 特徴エンコーダー
特徴エンコーダーは非常にシンプルで1次元畳み込みや正規化を組み合わせたものになっています。

```python
def conv_module(hidden_size, dropout=0.1):
    inputs = tf.keras.Sequential([
        tf.keras.layers.Conv1D(
            hidden_size, kernel_size=1, padding="same"),
        tf.keras.layers.LayerNormalization(
            epsilon=1e-6),
        tf.keras.layers.DepthwiseConv1D(
            kernel_size=3, padding="same"),
        tf.keras.layers.Dropout(dropout),
        tf.keras.layers.LayerNormalization(
            epsilon=1e-6),
    ])

    return inputs


class FeatureExtractorLayer(tf.keras.layers.Layer):
    def __init__(
        self,
        filter_sizes,
        kernel_sizes,
        strides,
        window_sizes,
        layer_id=0,
        **kwargs
    ):
        super().__init__(**kwargs)

        self.filter_sizes = filter_sizes
        self.kernel_sizes = kernel_sizes
        self.strides = strides
        self.window_sizes = window_sizes
        self.layer_id = layer_id

        conv_dim = filter_sizes[layer_id]
        kernel_size = kernel_sizes[layer_id]
        stride = strides[layer_id]
        self.conv_module = conv_module(conv_dim)

        self.conv_layer = tf.keras.layers.Conv1D(
            conv_dim,
            kernel_size,
            strides=stride,
            use_bias=False,
            padding="same",
            kernel_initializer="he_normal"
        )

    def call(self, inputs, training=False):
        inputs = self.conv_module(inputs, training=training)
        inputs = self.conv_layer(inputs)
        inputs = tf.keras.activations.relu(inputs)

        return inputs
```

### コンテキストネットワーク
transformerの出力と線形層（量子化層）の出力から対照損失を計算します。
ポジティブペアの類似度とネガティブペアの類似度の平均を表示するようにしています。
```python
class MaskedEncoder(tf.keras.layers.Layer):
    def __init__(
        self,
        hidden_size,
        num_heads,
        num_layers,
        intermediate_size,
        batch_size,
        patch_length,
        codebook_size,
        embedding_dim,
        num_quantizers,
        ema_decay,
        commitment_cost,
        threshold_ema_dead_code,
        sample_codebook_temperature,
        dropout=0.1,
        layer_norm_eps=1e-5,
        is_gelu_approx=False,
        norm_first=True,
        temperature=1.0,
        **kwargs
    ):
        super().__init__(**kwargs)

        self.hidden_size = hidden_size
        self.num_heads = num_heads
        self.num_layers = num_layers
        self.intermediate_size = intermediate_size
        self.dropout = dropout
        self.layer_norm_eps = layer_norm_eps
        self.is_gelu_approx = is_gelu_approx
        self.norm_first = norm_first
        self.patch_length = patch_length
        self.batch_size = batch_size
        self.temperature = temperature

        self.transformer = DilatedTransformer(
            hidden_size=self.hidden_size,
            num_heads=self.num_heads,
            intermediate_size=self.intermediate_size,
            num_layers=self.num_layers,
            layer_norm_eps=self.layer_norm_eps,
            is_gelu_approx=self.is_gelu_approx,
            dropout=self.dropout,
            norm_first=self.norm_first,
            attention_window_size=5,
            batch_size=batch_size,
            seq_len=patch_length
        )

        self.dense = tf.keras.layers.Dense(
            embedding_dim,
            dtype=tf.float32)

        self.final_projection = tf.keras.layers.Dense(
            embedding_dim, dtype=tf.float32)

        self.num_negative_samples = 100
        self.spec_augment = MixStripes(
            dim=1,
            mix_width=10,
            stripes_num=200,
            mask_ratio=0.065,
            random_noise_mask=True,
            fixed_stripes_num=False)

    def call(self, inputs, attention_mask=None, training=False, add_loss=True):
        quantized = self.dense(inputs)

        mask = None
        if add_loss:
            inputs, mask = self.spec_augment(
                inputs, training=training, add_loss=add_loss)

            mask = mask[:, :, 0]

        inputs = self.transformer(inputs, training=training)
        inputs = self.final_projection(inputs)

        if add_loss:
            contrastive_loss, pos_sim, neg_sim = self.contrastive_loss(
                mask, inputs, quantized)
            self.add_loss(contrastive_loss)
            self.add_metric(contrastive_loss, name="context")
            self.add_metric(pos_sim, name="pos_sim")
            self.add_metric(neg_sim, name="neg_sim")

        return inputs, quantized, mask

    def contrastive_loss(self, mask, context, quantized):
        loss = []
        pos_sims = []
        neg_sims = []
        for batch_index in range(self.batch_size):
            mask_indices = tf.where(mask[batch_index])
            c_t = tf.gather_nd(context[batch_index], mask_indices)
            q_t = tf.gather_nd(quantized[batch_index], mask_indices)

            pos_similarity = cosine_similarity(c_t, q_t)
            pos_sims.append(tf.reduce_mean(pos_similarity))

            neg_similarity = cosine_similarity_matmul(c_t, q_t)
            neg_sims.append(tf.reduce_mean(neg_similarity))

            neg_similarity_mask = tf.linalg.diag(
                tf.ones(tf.shape(neg_similarity)[-1]))
            neg_similarity = tf.exp(
                neg_similarity / self.temperature) * (1 - neg_similarity_mask)

            random_indices = tf.random.shuffle(
                tf.range(tf.shape(neg_similarity)[-1]))[:self.num_negative_samples]
            neg_similarity = tf.gather(neg_similarity, random_indices, axis=1)

            neg_similarity = tf.concat(
                [neg_similarity, pos_similarity[:, tf.newaxis]], axis=-1)

            numerator = tf.exp(pos_similarity / self.temperature)
            denominator = tf.reduce_sum(neg_similarity, axis=-1)
            loss_m = -tf.math.log(numerator / denominator)
            loss.append(tf.reduce_mean(loss_m))

        return tf.reduce_mean(loss), tf.reduce_mean(
            pos_sims), tf.reduce_mean(neg_sims)
```

### 量子化モデル
入力を量子化するためのモデルです。このモデルは単体で先に学習します。
```python
class QuantizeModel(tf.keras.Model):
    def __init__(
            self,
            config: MusicEncoderConfig,
            batch_size,
            seq_len,
            **kwargs):
        super().__init__(**kwargs)
        self.config = config
        self.batch_size = batch_size
        self.seq_len = seq_len

        self.residual_vq = ResidualVQ(
            input_dim=252 * 3,
            codebook_size=config.codebook_size,
            embedding_dim=config.quantizer_embedding_dim,
            num_quantizers=config.num_quantizers,
            batch_size=batch_size,
            ema_decay=config.ema_decay,
            threshold_ema_dead_code=config.threshold_ema_dead_code,
            commitment_cost=config.commitment_cost,
            sample_codebook_temperature=config.sample_codebook_temperature,
            kmeans_init=config.kmeans_init,
            dtype=tf.float32
        )

    def call(
            self,
            inputs,
            training=False):
        inputs = tf.transpose(inputs, (0, 1, 3, 2))
        inputs = tf.reshape(inputs, (self.batch_size, -1, 252 * 3))

        quantized_inputs, encoding_indices = self.residual_vq(
            inputs, training=training)
        quantized_inputs = tf.keras.activations.relu(quantized_inputs)

        quantized_loss = tf.reduce_mean(
            tf.square(
                tf.cast(
                    inputs,
                    dtype=tf.float32) - quantized_inputs))
        self.add_loss(quantized_loss)
        self.add_metric(quantized_loss, name="quantized")

        return quantized_inputs, encoding_indices
```

### 全体
全体は以下のようになります。入力はCQTスペクトログラムで12 * 3 * 7のL,R,L-Rのデータです。
```python
class MusicEncoder(tf.keras.Model):
    def __init__(
            self,
            config: MusicEncoderConfig,
            batch_size,
            seq_len,
            **kwargs):
        super().__init__(**kwargs)
        self.config = config
        self.batch_size = batch_size
        self.seq_len = seq_len

        self.feature_extract_layers = [
            FeatureExtractorLayer(
                filter_sizes=config.filter_sizes,
                kernel_sizes=config.kernel_sizes,
                strides=config.strides,
                window_sizes=config.window_sizes,
                is_gelu_approx=config.is_gelu_approx,
                layer_id=i
            )
            for i in range(len(config.filter_sizes))
        ]

        encoded_seq_len = seq_len

        self.masked_encoder = MaskedEncoder(
            hidden_size=config.hidden_size,
            num_heads=config.num_heads,
            num_layers=config.num_layers,
            intermediate_size=config.intermediate_size,
            batch_size=batch_size,
            patch_length=encoded_seq_len,
            codebook_size=config.codebook_size,
            embedding_dim=config.embedding_dim,
            num_quantizers=config.num_quantizers,
            ema_decay=config.ema_decay,
            commitment_cost=config.commitment_cost,
            threshold_ema_dead_code=config.threshold_ema_dead_code,
            sample_codebook_temperature=config.sample_codebook_temperature,
            dropout=config.dropout,
            layer_norm_eps=config.layer_norm_eps,
            is_gelu_approx=config.is_gelu_approx,
            norm_first=config.norm_first,
            temperature=config.temperature
        )

        model_input = tf.keras.layers.Input(shape=(None, 12 * 3 * 7, 3))
        self.quantize_layer = QuantizeModel(
            config, batch_size=batch_size, seq_len=seq_len)
        model_out = self.quantize_layer(model_input)
        self.quantize_model = tf.keras.Model(
            inputs=[model_input], outputs=model_out)

        self.projection_layers = [
            tf.keras.layers.Dense(
                config.quantizer_embedding_dim,
                dtype=tf.float32) for i in range(
                config.num_quantizers)]

    def call(
            self,
            inputs,
            training=False,
            add_loss=True):
        original_inputs = inputs
        inputs = tf.transpose(inputs, (0, 1, 3, 2))
        inputs = tf.reshape(inputs, (self.batch_size, -1, 252 * 3))

        quantized_inputs, encoding_indices = self.quantize_model(
            original_inputs, training=training)

        for i, feature_extractor_layer in enumerate(
                self.feature_extract_layers):
            inputs = feature_extractor_layer(inputs, training=training)
        feature = inputs

        inputs, quantized, mask = self.masked_encoder(
            inputs, training=training, add_loss=add_loss)

        embeddings = self.quantize_layer.residual_vq.get_embeddings()
        mlm_losses = []
        for i in range(len(self.projection_layers)):
            proj = self.projection_layers[i](inputs)[..., tf.newaxis, :]
            encoded = tf.transpose(embeddings[i])[tf.newaxis, tf.newaxis, ...]

            similarity = cosine_similarity_matmul(proj, encoded) / 0.1
            similarity = tf.squeeze(similarity, axis=-2)

            masked_similarity = tf.gather_nd(similarity, tf.where(mask))
            masked_labels = tf.gather_nd(encoding_indices[i], tf.where(mask))
            loss = tf.keras.losses.SparseCategoricalCrossentropy(
                from_logits=True)(masked_labels, masked_similarity)
            mlm_losses.append(loss)

        self.add_loss(tf.reduce_mean(mlm_losses))
        self.add_metric(
            tf.reduce_mean(mlm_losses),
            name="mlm_loss")

        return inputs

    def freeze_feature_extractor(self):
        for i in range(len(self.feature_extract_layers)):
            self.feature_extract_layers[i].trainable = False
```

### パラメータ
モデルのパラメータは以下のようになります。
```python
from dataclasses import dataclass, field

@dataclass
class MusicEncoderConfig:
    # feature extractor
    filter_sizes: list = field(
        default_factory=lambda: [512, ]
    )
    kernel_sizes: list = field(
        default_factory=lambda: [4, ])
    strides: list = field(default_factory=lambda: [1, ])

    # encoder
    hidden_size: int = 512
    num_layers: int = 11
    embedding_dim: int = 512
    temperature: float = 0.1
    dropout: float = 0.1
    num_heads: int = 8
    intermediate_size: int = 2048
    is_gelu_approx: bool = False
    layer_norm_eps: float = 1e-6
    norm_first: bool = True

    # quantizer
    quantizer_embedding_dim: int = 64
    codebook_size: int = 1024
    commitment_cost: float = 0.0
    num_quantizers: int = 8
    ema_decay: float = 0.99
    threshold_ema_dead_code: float = 2.0
    sample_codebook_temperature: float = 0.0
    kmeans_init: bool = False
```

# 学習
事前学習は約5000時間の音楽データを使用し、以下のような設定で学習しました。


```python
policy = tf.keras.mixed_precision.Policy('mixed_float16')
tf.keras.mixed_precision.set_global_policy(policy)
print('Compute dtype: %s' % policy.compute_dtype)
print('Variable dtype: %s' % policy.variable_dtype)

model_name = "music_encoder"

epochs = 200
batch_size = 2
accum_steps = 2
patch_len = 8192
log_dir = "./logs/music_encoder"

x_train, x_test, dataset = load_from_npz()
monitor = 'val_loss'
model_input = tf.keras.layers.Input(shape=(None, 12 * 3 * 7, 3))
config = MusicEncoderConfig()
music_encoder = MusicEncoder(
    config=config,
    batch_size=batch_size,
    seq_len=patch_len)
# 量子化モデルは事前に学習済み (重みは固定)
music_encoder.load_quantize_model(
        "./model/quantize_model/quantize_model.ckpt")
music_encoder_out = music_encoder(model_input)
model = tf.keras.Model(inputs=[model_input], outputs=music_encoder_out)

train_gen = DataGeneratorBatch(
    files=x_train,
    dataset=dataset,
    batch_size=batch_size,
    patch_length=patch_len,
    initialepoch=initial_epoch,
    max_queue=2,
    cache_size=cache_size)

test_gen = DataGeneratorBatch(
    files=x_test,
    dataset=dataset,
    batch_size=batch_size,
    patch_length=patch_len,
    validate_mode=True,
    cache_size=cache_size)

lrs = WarmupCosineDecay(
    len(train_gen) *
    epochs,
    warmup_steps=len(train_gen) *
    epochs *
    0.1,
    target_lr=1e-5)
callbacks = [
    lrs
]

model = GradientAccumulateModel(
    accum_steps=accum_steps,
    inputs=model.inputs,
    outputs=model.outputs,
    mixed_precision=True)
optimizer = tf.keras.optimizers.Adam(
    learning_rate=1e-3,
    beta_1=0.9,
    beta_2=0.98)

model.compile(optimizer=optimizer)
model.summary()

history = model.fit(
    x=train_gen,
    validation_data=test_gen,
    initial_epoch=initial_epoch,
    epochs=epochs,
    shuffle=False,
    callbacks=callbacks
)
```

OptimizerのAdamの設定ですが、[W2v-BERT](https://arxiv.org/pdf/2108.06209.pdf)の設定を参考にしました。beta2の値を0.98や0.95のような値にすると学習の安定性が上がるようです。[^7]
学習率は普通の学習と比べるとかなり低めの1e-5に設定しています。これは学習が非常に不安定で高い学習率を使用すると崩壊してしまうためです。accum_stepsの値を上げて実質のバッチサイズを上げても安定性は改善しませんでした。

入力シーケンスの長さは8192に設定しています。1フレームあたり0.02秒なので、0.02秒 × 8192 = 約163秒の長さになります。


### 学習結果

**対照損失**
![epoch_context (1).png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1275467/b6578e85-1290-392d-47f3-43ed0863662b.png)

**音響情報損失**
![epoch_mlm_loss.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1275467/303e620d-971a-bcb1-6776-ec44e08d23e1.png)

**ポジティブペアの類似度**
![epoch_pos_sim (1).png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1275467/00b245bb-01a1-65db-b10a-bede1c0ebb34.png)

**ネガティブペアの類似度**
![epoch_neg_sim (3).png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1275467/a704b1fe-90e3-58a3-4bbf-a51a2a216cfb.png)



# ファインチューニング
学習した事前学習モデルを使ってどの程度の効果が見られるか検証します。
比較するモデルは、LSTMを使ったモデル、事前学習ありのモデル、事前学習なしのモデル、音響情報損失関数なしのモデルの四つです。データ拡張は利用せず少量データでの効果を検証します。

ビート予測の評価は[Ballroom](https://github.com/CPJKU/BallroomAnnotations)データセットを使い、学習データとテストデータを7:3で分割します。ビート予測の後処理にはDBNBeatTrackingProcessorを使います。
最終的な評価はテストデータをmir_evalの**beat.f_measure**を使用して計算します。

コード予測の評価は独自のデータセット(50時間)を使い、学習データとテストデータを7:3で分割します。
評価はテストデータをcategorical_accuracyを使用して計算します。

ジャンル予測の評価は[GTZAN](https://www.kaggle.com/datasets/andradaolteanu/gtzan-dataset-music-genre-classification)データセットを使い、学習データとテストデータを7:3で分割します。
評価はテストデータをcategorical_accuracyを使用して計算します。

| モデル               | F1(Beat) | Acc(Chord) | Acc(Genre) |
|----------------------|----------|-----------|-----------|
| LSTM                 | 0.30     | 0.40 | 0.72 |
| with pre-training    | 0.7      | 0.51 | 0.72 |
| without pre-training | Failed   | 0.50 | 0.17 |
| without MLM          | 0.65     | 0.30 | 0.66 |

ビート予測の結果から、事前学習の効果が高いことがわかります。音響情報の損失関数を持つモデルも、持たないモデルに比べて同等あるいはそれ以上の精度を達成していることから、音響情報の損失関数を導入してもビートの情報が失われていないと言えます。
コードの予測結果では、音響情報の損失関数を持たないモデルはLSTMと比較して精度が大幅に低下しています。一方、音響情報の損失関数を持つモデルはLSTMよりも高い性能を示していて、音響情報の損失関数の導入がコードの精度向上に寄与していると考えられます。

以下は各モデルのコードの混同行列です。音響情報損失関数ありのモデルが一番正しく予測できていることが確認できます。

**音響情報損失関数あり**
![confusion_matrix_chord_net_with_mlm.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1275467/62aace96-0499-1a2a-83bc-7bea4ed369c1.png)

**音響情報損失関数なし**
![confusion_matrix_chord_net_without_mlm.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1275467/b6492637-367d-8532-e8f2-7a04ef4ef29b.png)

**LSTM**
![confusion_matrix_lstm.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1275467/b49df526-93da-624b-31b5-18d320386294.png)

**事前学習なし**
![confusion_matrix_chord_net_without_pretrain.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1275467/13bb64be-9321-e7dc-1fe8-e73f7cd3f5bc.png)



# まとめ
この記事では、wav2vec2.0を改造して音楽の特徴量を作成しました。wav2vec2.0の課題である音響情報少なさを解消するため、HuBERTの学習方法を参考に音響情報損失関数を導入しました。これにより、音響情報を必要とする下流タスクの精度が向上しました。

初めて事前学習の手法を試しましたが想定以上の効果があり驚きました。
まだ試せていないことが多くあるので時間があれば試していきたいと思います。


# おまけ
### コンテキストネットワークの中間層の出力
以下はコンテキストネットワークの各層の出力をアニメーションにしたものです。
層が深くなるにつれてマスク部分が補完されていく様子がわかると思います。
![中間層.gif](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1275467/4178c0b7-e8f8-f2ff-e4df-6f8ae130fff0.gif)


### 学習が正しく行えなかったセットアップ
**学習率を高めに設定**
wav2vec2.0の論文では学習率は5e-4や3e-4などの値を使用していますが、論文通りでは学習が崩壊してしまいました。
これはデータセットの大きさやバッチサイズなどが関係ありそうですが細かい検証はしていません。
論文によると60k時間のデータセットで事前学習した際は上記の学習率ですが、960時間の際は特徴エンコーダーの勾配を10分の1にしたようです。

**特徴エンコーダーに受容野が広いレイヤーを使う**
特徴エンコーダーにLSTMを使ったところ損失が急速に下がり特徴量に崩壊が見られました。
正常に学習できた特徴量はある程度の塊があるのですが、崩壊した際には完全にランダムに見える特徴量となっていました。
また、ポジティブペアの類似度は学習初期の時点で0.9近くに上昇していて、これは50%以上がマスクされることを考えると異常です。

### 自己類似度行列をプロットしてみる
以下は、コンテキストネットワークの出力の自己類似度行列です。
自己類似度行列は各時点のデータがほかの各時点のデータと類似しているかを視覚的に確認できます。
斜めの白い部分が類似している部分で、AメロやBメロ、もっと細かい単位での繰り返し要素が捉えられていることがわかります。
斜めよりも少し曲がっている部分がありますがこれは同じセクションでテンポが変化したことを表しています。
![cosine-min.jpg](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/1275467/58e74613-ac2a-f373-76cd-7872d8565632.jpeg)

### 今回検証できなかったこと
- 対照損失なしで音響情報の損失のみで学習 (音響情報の損失だけでビート情報を捉えられるのか？)
- 音響情報の損失用のtransformerを導入する (w2v-BERTのようなイメージ)
- ~~量子化を高精度にしたらどうなるか（RVQのレイヤー数を増やす、内部の次元数を上げるなど...)~~
- ~~特徴エンコーダーの大きさによる下流タスクへの影響~~
    - 特徴エンコーダーを大きくすると特徴エンコーダーがずるをするようになって下流タスクの精度が上がらなくなってしまいました


# 追記
MLMの音響情報なしで追加で検証してみました。M0は上記で検証したモデルです。
|モデル|時間解像度*1(frame)|マスクスパン(ms)|マスク率|量子化|Chord Acc|
|-----|-------------------|---------------|--------|-----|---------|
|M0|50|200|65|なし|0.3064|
|M1|12.5|400|32.5|RVQ(layer=8)|0.3797|
|M2|12.5|1600|50|RVQ(layer=8)|0.4373|
|M3|12.5|400|50|RVQ(layer=8)|0.4089|
|M4|12.5|2080|50|RVQ(layer=8)|0.4521|
|M5|12.5|2080|50|なし|0.4431|
|M6|12.5|2080|50|RVQ(layer=1)|0.2193|
|M7|12.5|2080|50|RVQ(layer=2)|0.2617|
|M8|12.5|2080|50|RVQ(layer=16)|**0.4628**|
|M9|12.5|4000|50|RVQ(layer=16)|0.4626|
|M10|12.5|2400|50|RVQ(layer=16)|0.4614|


***1: 1秒当たりのフレーム数**

### マスクスパン
M1、M2、M4の結果からマスクスパンが長いほどコードの精度が上がることがわかります。
マスクスパンを長くすることによってマスクの予測が難しくなり、より広いコンテキストを読み取ろうとするためだと思います。
マスクスパンが長いM4、M9、M10のモデルはアンダーフィットしていたように見え、さらに学習することでまだ精度が向上しそうでした。

### 量子化
量子化層はレイヤー数が増えるほど精度が向上するようです。M6やM7ではレイヤー数が少なく、表現力が低くなってしまい精度が低下していると思われます。量子化層を線形層に置き換えた場合ですがそれなりの精度は出ています。計算量を減らしたい場合は線形層でもありかもしれません。

## マスクスパンを変化させたときのキー予測精度
|モデル|時間解像度*1(frame)|マスクスパン(ms)|マスク率|量子化|Key Acc|
|-----|-------------------|---------------|--------|-----|---------|
|M8|12.5|2080|50|RVQ(layer=16)|0.5226|
|M9|12.5|4000|50|RVQ(layer=16)|**0.5767**|
|M10|12.5|2400|50|RVQ(layer=16)|0.5555|

マスクスパンを長くすることでキーの予測精度が向上しました。
キーの予測には長いコンテキストが有利なので、より長いコンテキストを捉えられる長いマスクスパンが有効だと思われます。ほかの長いコンテキストを必要とするようなタスクにも有効な可能性があります。(セクション予測など)

[^1]: InfoNCEと呼ばれる式です。
[^4]: https://arxiv.org/pdf/2010.10504v2.pdf
[^7]: https://engineering.linecorp.com/ja/blog/3.6-billion-parameter-japanese-language-model