Architecture and Design Backend Architecture Machine Learning Integration Beat Transformer Model

Beat-Transformer Model

Introduction

This document describes the Beat-Transformer deep learning model for simultaneous beat and downbeat tracking. It explains the Dilated Self-Attention architecture, model layers configuration, training methodology, spectrogram preprocessing pipeline, data loading mechanisms, audio demixing using Spleeter, checkpoint management, inference pipeline, performance optimization, eight-fold cross-validation training, ablation study implementations, and evaluation metrics. It also covers how the model handles different musical genres and tempos.

Project Structure

The Beat-Transformer implementation resides under the Beat-Transformer model directory. Key areas:

Model definition and layers: code/DilatedTransformer.py, code/DilatedTransformerLayer.py
Training and evaluation: code/train.py, code/eight_fold_test.py, code/utils.py
Data loading and preprocessing: code/spectrogram_dataset.py
Audio demixing and spectrogram generation: preprocessing/demixing.py, demix_spectrogram.py
Inference demos and wrappers: beat_tracking_demo.py, beat_tracking_fix.py
Ablation studies: code/ablation_models/*.py

graph TB
subgraph "Beat-Transformer"
A["DilatedTransformer.py"]
B["DilatedTransformerLayer.py"]
C["train.py"]
D["eight_fold_test.py"]
E["spectrogram_dataset.py"]
F["utils.py"]
G["demixing.py"]
H["demix_spectrogram.py"]
I["beat_tracking_demo.py"]
J["beat_tracking_fix.py"]
K["ablation_models/music_transformer.py"]
L["ablation_models/tcn.py"]
end
A --> B
C --> A
C --> E
C --> F
D --> A
D --> E
D --> F
E --> G
G --> H
I --> A
I --> F
J --> I
K --> B
L --> E

Diagram sources

Section sources

README.md:1-75

Core Components

DilatedTransformerModel: A convolutional front-end followed by stacked DilatedTransformerLayer blocks and two heads: one for beat/downbeat classification and another for tempo regression.
DilatedTransformerLayer: Implements dilated self-attention with relative positional embeddings across exponentially expanding receptive fields.
audioDataset and dataset_processing: Load and preprocess demixed spectrograms, quantize beats/downbeats, infer tempo histograms, and support 8-fold cross-validation splits.
Training loop: BCE loss for beat/downbeat, BCE loss for tempo, LR scheduling, and evaluation via DBN decoding.
Inference: Produces activation maps and optional attention accumulation for visualization.

Section sources

Architecture Overview

The Beat-Transformer combines CNN feature extraction with dilated self-attention to track beats and downbeats simultaneously. The model takes a 5-channel demixed spectrogram and predicts per-frame probabilities for beats and downbeats, plus a tempo distribution.

classDiagram
class Demixed_DilatedTransformerModel {
+int nhead
+int nlayers
+int attn_len
+int head_dim
+int dmodel
+forward(x)
+inference(x)
}
class DilatedTransformerLayer {
+forward(x, layer)
+inference(x, layer)
}
class DilatedMultiheadSelfAttentionWithRelativePositionalEmbedding {
+forward(query, key, value, layer)
+kv_roll(tensor, layer, ...)
}
Demixed_DilatedTransformerModel --> DilatedTransformerLayer : "uses"
DilatedTransformerLayer --> DilatedMultiheadSelfAttentionWithRelativePositionalEmbedding : "uses"

Diagram sources

Detailed Component Analysis

Dilated Self-Attention Mechanism

Relative positional embedding Er controls attention across dilated windows per layer.
kv_roll expands keys/values across exponentially increasing receptive fields controlled by 2^layer.
Attention is computed as scaled dot-product plus relative term, followed by softmax and value aggregation.
During inference, the attention matrices are accumulated to reconstruct full attention maps.

flowchart TD
Start(["Layer Input x"]) --> QKV["Compute Q, K, V"]
QKV --> RollK["Roll K across dilated windows<br/>kv_roll(layer)"]
QKV --> RollV["Roll V across dilated windows<br/>kv_roll(layer)"]
RollK --> QK["Compute QK"]
RollV --> Att["Add relative term Er<br/>Softmax over attention logits"]
QK --> Att
Att --> Out["Multiply by V and reshape"]
Out --> Skip["Residual + FFN"]
Skip --> End(["Layer Output"])

Diagram sources

Section sources

Model Layers Configuration

Convolutional front-end: 3 conv blocks with max pooling and ReLU to produce a compact representation.
Stacked DilatedTransformerLayer blocks: 9 layers with mixed attention modes—time-only dilated attention for early layers and interleaved temporal/instrument attention in middle layers.
Output heads:
- Beat/downbeat classification head (2 classes).
- Tempo head (300-bin categorical distribution).

sequenceDiagram
participant X as "Input Spectrogram<br/>(B, 5, T, F)"
participant CNN as "Conv Front-End"
participant Layers as "DilatedTransformer Layers"
participant Head1 as "Beat/Downbeat Head"
participant Head2 as "Tempo Head"
X->>CNN : Conv + Pool + ReLU
CNN->>Layers : Transpose to (B, T, D)
Layers-->>Head1 : Final representation
Layers-->>Head2 : Skip features + mean
Head1-->>Head1 : Linear + Sigmoid
Head2-->>Head2 : Linear + Softmax

Diagram sources

DilatedTransformer.py:41-90

Section sources

DilatedTransformer.py:7-40

Training Methodology

Eight-fold cross-validation: audioDataset folds data across datasets; each fold trains on 7/8 and validates/test on 1/8.
Losses:
- Binary cross-entropy for beat/downbeat with masked invalid positions.
- Categorical cross-entropy for tempo distribution.
Optimizer and scheduler: RAdam with Lookahead, ReduceLROnPlateau.
Evaluation: DBN decoding for beat and downbeat; metrics include f-measure, cmlt, amlt.

sequenceDiagram
participant Loader as "DataLoader"
participant Model as "Demixed_DilatedTransformerModel"
participant Opt as "Optimizer"
participant Sch as "Scheduler"
Loader->>Model : Batch (data, beat_gt, downbeat_gt, tempo_gt)
Model->>Model : forward()
Model-->>Opt : loss = BCE + BCE_tempo
Opt->>Opt : backward() and step()
Sch->>Sch : step(val_loss)

Diagram sources

Section sources

Spectrogram Preprocessing Pipeline

Demixing: Uses Spleeter 5-stems separation; optionally falls back to simple Mel spectrogram.
Mel filterbank: 128 mel bins, hop length matching model fps (44100/1024).
Data preparation: Aggregates multiple datasets, saves compressed npz files for spectrograms and annotations.

flowchart TD
A["Raw Audio"] --> B["Spleeter 5-Stems Separation"]
B --> C["Mel Spectrogram per Stem<br/>128 bins, hop=1024"]
C --> D["Stack (5, T, F)"]
D --> E["Save demix_spectrogram_data.npz"]
A --> F["Fallback: Simple Mel Spectrogram"]
F --> D

Diagram sources

Section sources

Data Loading Mechanisms

audioDataset loads preprocessed spectrograms and annotations, splits into folds, and supports train/validation/test splits.
dataset_processing quantizes events to model fps, applies smoothing, infers tempo histograms, and clips sequences to fixed lengths.

flowchart TD
A["Load npz spectrograms & annotations"] --> B["Shuffle and split into 8 folds"]
B --> C["Train: Clip sequences and augment"]
B --> D["Val/Test: Clip to ~420s segments"]
C --> E["dataset_processing: quantize beats/downbeats<br/>smooth & infer tempo"]
D --> E
E --> F["Return (instr, time, mel) batches"]

Diagram sources

Section sources

spectrogram_dataset.py:17-282

Eight-Fold Cross-Validation and Testing

eight_fold_test.py runs inference across all 8 folds, aggregates activations, decodes with DBN, and computes metrics per dataset.

sequenceDiagram
participant Test as "eight_fold_test.py"
participant DS as "audioDataset.get_fold()"
participant Model as "Demixed_DilatedTransformerModel"
participant Eval as "utils metrics"
loop for Fold in 0..7
Test->>DS : get_fold(Fold)
DS-->>Test : train/val/test sets
Test->>Model : load checkpoint Fold
Test->>Model : forward on test loader
Model-->>Test : activations
Test->>Eval : DBN decoding + metrics
end

Diagram sources

Section sources

Ablation Study Implementations

music_transformer.py: Standard transformer encoder with relative positional embeddings and dilation-aware masking.
tcn.py: Temporal convolution network with dilated convolutions for comparison. These modules demonstrate alternatives to the dilated self-attention architecture.

Section sources

Evaluation Metrics

Beat/Downbeat decoding via DBN with configurable thresholds and tempo priors.
Metrics include f-measure, cmlt, amlt averaged across songs.

Section sources

Inference Pipeline and Attention Visualization

beat_tracking_demo.py runs the model on demixed spectrograms, supports chunked inference for long audio, and decodes beats/downbeats with DBN.
Optional attention accumulation during inference enables visualization of effective receptive fields.

sequenceDiagram
participant Spec as "Demixed Spec"
participant Model as "Demixed_DilatedTransformerModel"
participant Act as "Activations"
participant DBN as "DBN Decoder"
Spec->>Model : forward()
Model-->>Act : beat/downbeat activations
Act->>DBN : decode beats/downbeats
DBN-->>DBN : compute metrics

Diagram sources

Section sources

Dependency Analysis

Internal dependencies:
- train.py depends on DilatedTransformer, spectrogram_dataset, utils.
- eight_fold_test.py depends on DilatedTransformer, spectrogram_dataset, utils.
- beat_tracking_demo.py depends on DilatedTransformer and utils.
- preprocessing/demixing.py depends on spleeter and librosa.
External libraries: torch, librosa, madmom, spleeter, numpy, scipy.

graph LR
Train["train.py"] --> DT["DilatedTransformer.py"]
Train --> DS["spectrogram_dataset.py"]
Train --> UT["utils.py"]
Eval["eight_fold_test.py"] --> DT
Eval --> DS
Eval --> UT
Demo["beat_tracking_demo.py"] --> DT
Demo --> UT
Prep["demixing.py"] --> DS
DS --> Prep

Diagram sources

Section sources

Performance Considerations

Model limitations: The model’s effective receptive field grows with layers and attention length; very long audio requires chunking with overlap blending.
Padding strategy: Adds context frames at the beginning to enable accurate early predictions.
Batch size and device: Single-batch training/evaluation to reduce memory pressure; GPU recommended for inference.
Data augmentation: Random mixing of stems during training to improve generalization.
Gradient clipping and LR scheduling: Reduce divergence and stabilize convergence.

[No sources needed since this section provides general guidance]

Troubleshooting Guide

Spleeter not available: demix_spectrogram.py falls back to simple Mel spectrogram generation.
Long audio handling: beat_tracking_demo.py processes audio in overlapping chunks and blends outputs.
Missing DBN detections: beat_tracking_demo.py includes librosa fallbacks and conservative defaults.
CUDA issues: beat_tracking_demo.py automatically falls back to CPU if CUDA is unavailable.

Section sources

Conclusion

The Beat-Transformer leverages dilated self-attention to jointly model beats and downbeats with strong generalization across genres and tempos. Its training methodology integrates robust data preprocessing, cross-validation, and DBN decoding for evaluation. The provided preprocessing, training, and inference utilities enable reproducible experiments and practical deployment.

[No sources needed since this section summarizes without analyzing specific files]

Appendices

Model Checkpoint Management

Training saves checkpoints per epoch with optimizer and scheduler states.
eight_fold_test.py loads fold-specific checkpoints and evaluates on test sets.

Section sources

Handling Different Genres and Tempos

Datasets: Ballroom, Carnatic, GTZAN, Hainsworth, SMC, Harmonix.
Tempo inference: Histogram-based tempo estimation from beat intervals with smoothing and interpolation.
Evaluation: Metrics computed per dataset and aggregated.

Section sources

ChordMiniApp Wiki

Home

General

API Reference