CLAUDE.md

This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.

Project Overview

This is the BEATs (Audio Pre-Training with Acoustic Tokenizers) project, part of Microsoft's UniLM family. It implements a self-supervised audio pre-training framework based on iterative acoustic tokenization and masked audio modeling. Paper: arXiv:2212.09058.

The repo also contains an active extension, Spatial-BEATs, which adds spatial audio understanding (direction-of-arrival, distance estimation) on top of the frozen BEATs encoder for First Order Ambisonics (FOA) data.

Key Dependencies

PyTorch, torchaudio (for fbank feature extraction via torchaudio.compliance.kaldi)
einops (used by quantizer for codebook k-means init)
Training uses torchrun for distributed data parallel

Training Commands

Spatial-BEATs (three-stage mono-AST on ov1 FOA data)

# All knobs overridable via env vars: GPUS, BATCH_SIZE, NUM_WORKERS, etc.
./run_ov1_ast_three_stage.sh

Stages: (1) class warmup with frozen BEATs, (2) spatial-first, (3) balanced classification + spatial.

Pre-trunk AST experiment (two-stage)

./run_ov1_pretrunk_ast_experiment.sh

Stages: (1) class-only warmup with task tokens inside BEATs trunk, (2) spatial CE finetune.

Single training run

torchrun --nproc_per_node=4 train_spatial_beats.py \
  --preset <preset_name> \
  --output-dir <output_dir> \
  --batch-size 8 --num-workers 4 --num-epochs 12

Available presets are defined via make_*_config() factories in train_spatial_beats.py and listed in spatial_beats_ov123_stage1_config.py.

Architecture

Original BEATs (inference-only weights)

Raw waveform (16kHz)
  → fbank (128 mel bins, frame_length=25ms, frame_shift=10ms)
  → normalize with fixed mean/std
  → Conv2d patch embedding
  → LayerNorm → optional Linear projection
  → TransformerEncoder (N layers with relative position bias + GRU gating)
  → extract_features() returns [B, T, D] representations
  → (finetuned models) → Linear predictor → sigmoid → class probabilities

Two model classes share this backbone:

BEATs (BEATs.py): audio encoder. extract_features() returns representations or class probs (if finetuned).
Tokenizers (Tokenizers.py): same encoder + NormEMAVectorQuantizer head. extract_labels() returns discrete codebook indices.

Spatial-BEATs extension

Builds on top of BEATs to add spatial audio capabilities:

SpatialBEATs (spatial_beats.py): wraps a frozen BEATs TransformerEncoder with multi-channel FOA preprocessing (SpatialBEATsPreprocessor), a SpatialPatchEmbedding for the extra channels, and task-specific prediction heads.
spatial_modules.py: contains all building blocks — SpatialPatchEmbedding, SpatialDeltaPatchAdapter, FixedSlotReadout, MonoTaskTokenReadout, FrequencyPool, TemporalResampler, and prediction heads (SpatialPredictionHeads, MonoTaskPredictionHeads, PreTrunkASTPredictionHeads).
spatial_dataset.py: SpatialDataset loads FOA audio from JSONL manifests with per-frame source annotations (azimuth, elevation, distance, class). Uses a Qwen-2.5-Omni-aligned mel frontend (16kHz, 128 bins, hop=160).
spatial_loss.py: multi-task loss with Hungarian-style slot matching — activity BCE, azimuth/elevation CE over binned angles, distance regression, and auxiliary source classification.

Module dependency graph

modules.py          — primitives: GradMultiply, SamePad, GLU_Linear, quant_noise, activation fns
quantizer.py        — NormEMAVectorQuantizer, EmbeddingEMA (VQ-VAE codebook with EMA updates)
backbone.py         — TransformerEncoder, TransformerSentenceEncoderLayer, MultiheadAttention
BEATs.py            — BEATs model (uses backbone)
Tokenizers.py       — Tokenizers model (uses backbone + quantizer)
spatial_modules.py  — spatial building blocks (patch embeddings, readout heads, prediction heads)
spatial_beats.py    — SpatialBEATs model (uses backbone + spatial_modules)
spatial_dataset.py  — SpatialDataset + collation
spatial_loss.py     — loss computation + slot matching (uses spatial_modules output types)
train_spatial_beats.py — training loop, presets, CLI (uses spatial_beats, spatial_dataset, spatial_loss)

Loading Pre-trained Checkpoints

Checkpoints are dict with keys 'cfg' (config dict) and 'model' (state dict):

checkpoint = torch.load('model.pt')
cfg = BEATsConfig(checkpoint['cfg'])
model = BEATs(cfg)
model.load_state_dict(checkpoint['model'])

Same pattern for Tokenizers with TokenizersConfig.

Audio Input Contract

All models expect 16kHz mono waveforms
preprocess() converts to 128-bin fbank features normalized with fixed mean=15.41663, std=6.55582
Padding masks are bool tensors where True = padded position
Spatial-BEATs uses 4-channel FOA input instead of mono

我希望在原始BEATs的基础上更改模型的框架，让模型有FOA音频的理解能力，能够在声源分类之外拥有识别位置的能力，这样的encoder作为我未来输入给LLM的例子。我之前自己尝试了一些做法，不过class分类不是很收敛，空间指标比如dis,ele,azimuth的loss几乎不收敛，我感觉我的方法太过于ML了，没有充分的利用DL的能力，或许应该一定程度上相信attention的能力来学习。我认为应该像BAT一样，你看这个目录下面的Spatial-AST的训练是从AudioMAE的训练开始的，我觉得确实应该学习他的设计来类似的训练我的Spatial-BEATs，我设计了实验run_ov1_pretrunk_ast_experiment.sh来验证，现在有了初步的结果，但是看的出来，还不是很收敛，预期结果和我想的完全不一样，我到底应该怎么办呢？还有疑问是BEATS是用audioset训练的，我现在的ov1数据干声来源于FSD50K，这是不是首先会影响分类任务，我是不是应该先在分类任务上finetune到一定的程度之后再考虑空间呢