muscriptor/models/btc.py

"""BTC: the Bi-directional Transformer for Chord Recognition (ISMIR 2019).

Inference-only port of Jonggwon Park's reference implementation
(https://github.com/jayg996/BTC-ISMIR19, MIT licensed, see LICENSE-BTC), kept
faithful enough to load the published checkpoints unchanged. Everything the
checkpoint has a weight for is reproduced exactly, quirks included:

* the layer norm divides by an *unbiased* standard deviation (`torch.std`),
  not the biased one `nn.LayerNorm` uses, so it can't be swapped for the
  built-in module;
* the position-wise feed-forward applies its ReLU after *every* convolution,
  the last one included, so its output is non-negative;
* attention is masked in both directions — one block sees only the past, its
  twin only the future — which is what makes the model "bi-directional".

Dropout is left out (this only ever runs under `eval()`), and so is the unused
LSTM the reference output layer carries: see `load_state_dict` below.
"""

import math

import numpy as np
import torch
import torch.nn.functional as F
from torch import nn


def timing_signal(length: int, channels: int) -> torch.Tensor:
    """Sinusoidal position encoding, [1, length, channels].

    The transformer's original formulation, as ported into BTC from
    tensor2tensor: `channels // 2` sine components followed by the same number
    of cosines, over timescales spanning 1 … 1e4.
    """
    position = np.arange(length)
    num_timescales = channels // 2
    log_increment = math.log(1.0e4) / (num_timescales - 1)
    inv_timescales = np.exp(
        np.arange(num_timescales, dtype=np.float64) * -log_increment
    )
    scaled = np.expand_dims(position, 1) * np.expand_dims(inv_timescales, 0)
    signal = np.concatenate([np.sin(scaled), np.cos(scaled)], axis=1)
    signal = np.pad(signal, [[0, 0], [0, channels % 2]], "constant")
    return torch.from_numpy(signal.reshape([1, length, channels])).float()


def causal_mask(length: int) -> torch.Tensor:
    """Additive attention mask, [1, 1, length, length], hiding future frames.

    Transposing the last two dimensions turns it into the mask that hides the
    past instead — which is exactly how the backward block is built.
    """
    mask = np.triu(np.full([length, length], -np.inf), 1)
    return torch.from_numpy(mask).float().unsqueeze(0).unsqueeze(1)


class LayerNorm(nn.Module):
    """Layer norm over the last dimension, with an unbiased standard deviation.

    `nn.LayerNorm` normalizes by the biased (population) standard deviation.
    BTC was trained against this variant, so the two are not interchangeable:
    swapping it in shifts every activation slightly and costs accuracy.
    """

    def __init__(self, features: int, eps: float = 1e-6):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(features))
        self.beta = nn.Parameter(torch.zeros(features))
        self.eps = eps

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.gamma * (x - mean) / (std + self.eps) + self.beta


class MultiHeadAttention(nn.Module):
    """Scaled dot-product attention over `num_heads` heads, biased by a mask."""

    def __init__(
        self,
        depth: int,
        key_depth: int,
        value_depth: int,
        num_heads: int,
        bias_mask: torch.Tensor,
    ):
        super().__init__()
        self.num_heads = num_heads
        self.query_scale = (key_depth // num_heads) ** -0.5
        # A buffer, not a plain attribute: it has to follow the model onto the
        # GPU, and it must stay out of the state dict (the checkpoint has no
        # entry for it) — hence persistent=False.
        self.register_buffer("bias_mask", bias_mask, persistent=False)
        self.query_linear = nn.Linear(depth, key_depth, bias=False)
        self.key_linear = nn.Linear(depth, key_depth, bias=False)
        self.value_linear = nn.Linear(depth, value_depth, bias=False)
        self.output_linear = nn.Linear(value_depth, depth, bias=False)

    def _split(self, x: torch.Tensor) -> torch.Tensor:
        batch, length, depth = x.shape
        return x.view(batch, length, self.num_heads, depth // self.num_heads).permute(
            0, 2, 1, 3
        )

    def _merge(self, x: torch.Tensor) -> torch.Tensor:
        batch, heads, length, depth = x.shape
        return x.permute(0, 2, 1, 3).contiguous().view(batch, length, depth * heads)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        queries = self._split(self.query_linear(x)) * self.query_scale
        keys = self._split(self.key_linear(x))
        values = self._split(self.value_linear(x))
        logits = torch.matmul(queries, keys.permute(0, 1, 3, 2))
        length = logits.shape[-1]
        logits = logits + self.bias_mask[:, :, :length, :length].type_as(logits)
        weights = F.softmax(logits, dim=-1)
        return self.output_linear(self._merge(torch.matmul(weights, values)))


class Conv(nn.Module):
    """1-D convolution over a [batch, length, channels] sequence, centre-padded."""

    def __init__(self, in_channels: int, out_channels: int, kernel_size: int = 3):
        super().__init__()
        self.pad = nn.ConstantPad1d((kernel_size // 2, (kernel_size - 1) // 2), 0)
        self.conv = nn.Conv1d(in_channels, out_channels, kernel_size=kernel_size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.conv(self.pad(x.permute(0, 2, 1))).permute(0, 2, 1)


class PositionwiseFeedForward(nn.Module):
    """BTC's convolutional feed-forward: conv → ReLU → conv → ReLU.

    The trailing ReLU is not a typo — the reference loops over its layers and
    activates after each one, so the block's output is non-negative. Removing
    it changes what the residual stream carries.
    """

    def __init__(self, hidden_size: int, filter_size: int):
        super().__init__()
        self.layers = nn.ModuleList(
            [Conv(hidden_size, filter_size), Conv(filter_size, hidden_size)]
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        for layer in self.layers:
            x = F.relu(layer(x))
        return x


class SelfAttentionBlock(nn.Module):
    """Pre-norm attention + feed-forward, each wrapped in a residual."""

    def __init__(
        self,
        hidden_size: int,
        key_depth: int,
        value_depth: int,
        filter_size: int,
        num_heads: int,
        bias_mask: torch.Tensor,
    ):
        super().__init__()
        self.multi_head_attention = MultiHeadAttention(
            hidden_size, key_depth, value_depth, num_heads, bias_mask
        )
        self.positionwise_convolution = PositionwiseFeedForward(
            hidden_size, filter_size
        )
        self.layer_norm_mha = LayerNorm(hidden_size)
        self.layer_norm_ffn = LayerNorm(hidden_size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = x + self.multi_head_attention(self.layer_norm_mha(x))
        return x + self.positionwise_convolution(self.layer_norm_ffn(x))


class BiDirectionalSelfAttention(nn.Module):
    """One BTC layer: a past-only and a future-only block, then a projection."""

    def __init__(
        self,
        hidden_size: int,
        key_depth: int,
        value_depth: int,
        filter_size: int,
        num_heads: int,
        max_length: int,
    ):
        super().__init__()
        forward_mask = causal_mask(max_length)
        params = (hidden_size, key_depth, value_depth, filter_size, num_heads)
        self.attn_block = SelfAttentionBlock(*params, forward_mask)
        self.backward_attn_block = SelfAttentionBlock(
            *params, torch.transpose(forward_mask, 2, 3)
        )
        self.linear = nn.Linear(hidden_size * 2, hidden_size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        both = torch.cat((self.attn_block(x), self.backward_attn_block(x)), dim=2)
        return self.linear(both)


class BiDirectionalSelfAttentionLayers(nn.Module):
    """The encoder stack: project the features in, add timing, run the layers."""

    def __init__(
        self,
        feature_size: int,
        hidden_size: int,
        num_layers: int,
        num_heads: int,
        key_depth: int,
        value_depth: int,
        filter_size: int,
        max_length: int,
    ):
        super().__init__()
        self.register_buffer(
            "timing_signal", timing_signal(max_length, hidden_size), persistent=False
        )
        self.embedding_proj = nn.Linear(feature_size, hidden_size, bias=False)
        self.self_attn_layers = nn.ModuleList(
            BiDirectionalSelfAttention(
                hidden_size,
                key_depth,
                value_depth,
                filter_size,
                num_heads,
                max_length,
            )
            for _ in range(num_layers)
        )
        self.layer_norm = LayerNorm(hidden_size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.embedding_proj(x)
        x = x + self.timing_signal[:, : x.shape[1], :].type_as(x)
        for layer in self.self_attn_layers:
            x = layer(x)
        return self.layer_norm(x)


class OutputLayer(nn.Module):
    """Linear projection from the encoder's hidden size to the chord vocabulary."""

    def __init__(self, hidden_size: int, output_size: int):
        super().__init__()
        self.output_projection = nn.Linear(hidden_size, output_size)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.output_projection(x)


# Architecture of the published checkpoints (run_config.yaml in the reference
# repo). `timestep` is both the attention mask's size and the number of frames
# the model is fed at a time, so it is a hard constraint on the caller, not a
# tunable: muscriptor.utils.chords sizes its feature blocks around it.
FEATURE_SIZE = 144
TIMESTEP = 108
HIDDEN_SIZE = 128
NUM_LAYERS = 8
NUM_HEADS = 4
FILTER_SIZE = 128
# The "large vocabulary" head: 12 roots x 14 qualities, plus X (unknown) and N.
NUM_CHORDS_LARGE_VOCA = 170


class BTCModel(nn.Module):
    """The published BTC network, minus everything only training needed."""

    def __init__(self, num_chords: int = NUM_CHORDS_LARGE_VOCA):
        super().__init__()
        self.self_attn_layers = BiDirectionalSelfAttentionLayers(
            FEATURE_SIZE,
            HIDDEN_SIZE,
            NUM_LAYERS,
            NUM_HEADS,
            HIDDEN_SIZE,
            HIDDEN_SIZE,
            FILTER_SIZE,
            TIMESTEP,
        )
        self.output_layer = OutputLayer(HIDDEN_SIZE, num_chords)

    def forward(self, features: torch.Tensor) -> torch.Tensor:
        """Chord logits for `features` [batch, TIMESTEP, FEATURE_SIZE]."""
        return self.output_layer(self.self_attn_layers(features))

    def load_published_state_dict(self, state_dict: dict) -> None:
        """Load a reference checkpoint's `model` dict into this network.

        The reference output layer allocates a bidirectional LSTM that its
        forward pass never calls, so the checkpoint carries weights this port
        has no module for. They are dropped here rather than loaded with
        `strict=False`, which would also swallow a genuinely mismatched
        checkpoint.
        """
        wanted = {k: v for k, v in state_dict.items() if "output_layer.lstm." not in k}
        self.load_state_dict(wanted)