docs/spatial_beats_training_overview.md

Spatial-BEATs Training And Architecture Overview

This document summarizes the current Spatial-BEATs implementation in this repository:

• model architecture
• tensor shape flow
• dataset contract
• variable-length batching
• supervision and losses
• stage-1 training setup
• current ov1/ov2/ov3 presets

The implementation described here corresponds to:

• spatial_beats.py
• spatial_modules.py
• spatial_dataset.py
• spatial_loss.py
• train_spatial_beats.py
• spatial_beats_ov123_stage1_config.py

1. Goal

Spatial-BEATs is a separate spatial encoder for FOA audio.

It is designed to:

• reuse the BEATs backbone and pretrained weights
• take full FOA input instead of only the W channel
• learn spatial structure through explicit supervision
• output fixed-rate spatial tokens for an LLM
• stay separate from the original audio encoder used for semantic audio understanding

The current implementation follows the simplified design:

• the main objective is to train the FOA front-end and BEATs trunk to produce spatially informative embeddings
• the supervision heads are lightweight readout heads
• the final LLM tokens are taken from the encoder-side spatial embeddings, not from the final logits

2. High-Level Architecture

The end-to-end model path is:

FOA waveform
  -> FOA spatial preprocessor
  -> multi-channel patch embedding
  -> BEATs trunk
  -> frequency pooling
  -> temporal resampling to 2.5 Hz
  -> shallow temporal readout
  -> spatial embeddings
  -> fixed-slot supervision heads
  -> projector
  -> LLM spatial tokens

More concretely:

[B, 4, T]
  -> [B, 7, T_f, 128]
  -> [B, N_p, 512]
  -> [B, N_p, 768]
  -> [B, T_p, 768]
  -> [B, T_s_max, 768]
  -> [B, T_s_max, 768]
  -> [B, T_s_max, 4, 768]
  -> [B, T_s_max, d_llm]

Where:

• B: batch size
• T: waveform length in samples
• T_f: acoustic frame count before patching
• N_p: number of BEATs patches
• T_p: time-axis patch count after frequency pooling
• T_s_max: padded token count in the batch after resampling to 2.5 Hz
• d_llm: spatial token width sent to the LLM

3. Input And Front-End

3.1 Input audio

The model expects:

• FOA waveform
• shape [B, 4, T]
• channel order: W, X, Y, Z
• sample rate: 16 kHz

3.2 Qwen-like low-level mel setup

The current front-end is aligned to the Qwen-2.5-Omni audio tower style low-level parameters:

• sample_rate = 16000
• num_mel_bins = 128
• n_fft = 400
• win_length = 400
• hop_length = 160
• dither = 0.0

These parameters are shared between:

• SpatialBEATsConfig
• SpatialDatasetConfig

This keeps the data pipeline and the model front-end consistent.

3.3 FOA feature construction

The preprocessor converts FOA waveform into a 7-channel feature map:

• W_logmel
• X_logmel
• Y_logmel
• Z_logmel
• IVx
• IVy
• IVz

Output shape:

• foa_feat: [B, 7, T_f, 128]

This allows the whole FOA structure to enter the backbone instead of relying on only W.

4. Backbone And Spatial Embedding Path

4.1 Spatial patch embedding

The model replaces the original single-channel patch stem with a 7-channel patch embedding:

• input: foa_feat [B, 7, T_f, 128]
• output: patch_tokens [B, N_p, 512]
• also returns grid_size = (T_p, F_p)

This is the first modified entry point for reusing BEATs on FOA input.

4.2 Reused BEATs trunk

The trunk reuses BEATs pretrained components:

• layer_norm
• post_extract_proj
• encoder.pos_conv
• all transformer layers
• encoder.layer_norm

Flow:

• input: patch_tokens [B, N_p, 512]
• output: encoder_memory [B, N_p, 768]

4.3 Frequency pooling

The patch sequence is reshaped back into a patch grid and pooled over the frequency axis:

• input: encoder_memory [B, N_p, 768] with grid_size=(T_p, F_p)
• reshaped internally to [B, T_p, F_p, 768]
• pooled output: temporal_patch_tokens [B, T_p, 768]

This produces a time-aligned sequence before the final token-rate conversion.

4.4 Temporal resampling

The temporal resampler converts the patch-rate sequence into the final spatial token rate:

• target token rate: 2.5 Hz
• per-sample target length:

T_s_i = round(duration_i * 2.5)

Batch handling:

• each sample is resampled independently
• the batch is padded to T_s_max = max_i(T_s_i)
• a temporal mask is produced

Outputs:

• temporal_tokens: [B, T_s_max, 768]
• temporal_padding_mask: [B, T_s_max]

Mask convention:

• False: valid time step
• True: padded time step

4.5 Shallow temporal readout

The shallow temporal readout refines the resampled sequence with a lightweight transformer encoder:

• input: temporal_tokens [B, T_s_max, 768]
• output: spatial_embeddings [B, T_s_max, 768]

This is the main representation used for both:

• spatial supervision
• final projection to LLM tokens

5. Supervision Heads

The current stage-1 design does not use a heavy decoder.

Instead, it uses a fixed-slot readout for supervision only.

5.1 Fixed-slot readout

The readout expands each time step into a small number of internal supervision slots:

• max slots per step: K = 4
• input: spatial_embeddings [B, T_s_max, 768]
• output: slot_latents [B, T_s_max, 4, 768]

Important:

• K=4 is only a supervision capacity
• it does not change the final LLM token count
• the final LLM-visible token rate is still 2.5 Hz

5.2 Prediction heads

Each supervision slot predicts:

• pred_activity: [B, T_s_max, 4]
• pred_azi_logits: [B, T_s_max, 4, 360]
• pred_ele_logits: [B, T_s_max, 4, 180]
• pred_dist: [B, T_s_max, 4, 1]
• pred_class_logits: [B, T_s_max, 4, C]

Where:

• C = 65
• the class vocabulary comes from:
  • /apdcephfs_cq12/share_302080740/user/schmittzhu/data/fsd50k/FSD50K.ground_truth/final_vocabulary.csv

These heads are used to supply explicit training loss and push the front-end plus BEATs trunk to learn spatial structure.

6. LLM Spatial Tokens

The final LLM tokens are not taken from slot logits.

They are projected from the encoder-side spatial embeddings:

• input: spatial_embeddings [B, T_s_max, 768]
• output: llm_spatial_tokens [B, T_s_max, d_llm]

Therefore:

• 2.5 Hz means final LLM-visible tokens arrive at 2.5 tokens/second
• a 20 s clip produces about 50 spatial tokens
• a 10 s clip produces about 25 spatial tokens

This is the externally visible spatial token interface.

7. Pretrained Weight Reuse

The model initializes from BEATs_iter3+ AS2M.

Current pretrained loading logic:

• selectively load BEATs trunk modules
• skip task-specific components that do not match
• inflate the old single-channel patch embedding into the new 7-channel stem

Patch stem initialization rule:

• original BEATs patch weight is copied into channel 0 of the new 7-channel stem
• remaining channels start from zero

This is a conservative initialization intended to preserve BEATs trunk stability while enabling FOA adaptation.

8. Dataset Contract

8.1 Supported manifests

The dataset loader currently supports:

• ov1_foa.jsonl
• ov2_foa.jsonl
• ov3_foa.jsonl

It handles:

• single-source top-level manifest style
• nested multi-source manifest style with sources

8.2 Required scene-level data

At scene level the dataset expects one FOA path, typically:

• output_foa_path

or compatible fallback names already handled in the parser.

8.3 Required source-level data

For each source, the loader extracts:

• source class
• azimuth
• elevation
• distance
• weak time window

Internally each source is converted into a SourceEvent containing:

• class_index
• class_label
• azimuth_deg
• elevation_deg
• distance_m
• start_time_seconds
• end_time_seconds

8.4 Vocabulary mapping

Source labels are mapped to final_vocabulary.csv.

The loader supports several field aliases, including:

• mono_target_label
• mono_primary_label
• final_label
• source_label
• label

and several id-style aliases if an integer class index is already present.

9. Variable-Length Batching

Handling mixed-length FOA clips is a core part of the current implementation.

9.1 Waveform padding

At batch time:

• each waveform is padded to the batch maximum waveform length
• the padded tensor has shape [B, 4, T_max]
• a waveform padding mask is created:

waveform_padding_mask: [B, T_max]

Mask convention:

• False: valid waveform sample
• True: padded sample

dui

9.2 Temporal token padding

After temporal resampling:

• each sample has its own T_s_i = round(duration_i * 2.5)
• the batch is padded to T_s_max
• the model returns:

temporal_padding_mask: [B, T_s_max]
target_num_steps: [B]

All temporal supervision, matching, and loss computation respect these lengths.

9.3 Long clip truncation

The current training presets cap clip duration at:

• 20.0 seconds

The dataset applies cropping before batching.

Preset crop policy:

• crop_mode = "start"

This means:

• clips longer than 20 seconds are truncated from the beginning
• training and validation follow the same deterministic sequence policy

If needed later, the dataset also supports:

• random
• center
• none

10. Matching And Losses

10.1 Weak temporal supervision

The model uses weak source windows:

• each source provides start_time_seconds and end_time_seconds
• these define a valid supervision window, not guaranteed frame-level activity

The loss code first converts source windows into a time-window mask:

window_mask: [B, N_gt, T_s_max]

10.2 Per-step fixed-slot matching

Matching is performed per time step:

• only on valid temporal positions
• only within each source's weak time window
• between active GT sources and the K=4 slot predictions

The current matcher uses a detached cost built from:

• activity
• class
• azimuth
• elevation
• distance

The output contains the assigned GT target for each valid slot-time pair.

10.3 Multi-task loss terms

Current loss terms are:

• loss_activity
• loss_azi
• loss_ele
• loss_dist
• loss_cls_aux
• loss_temp

Their roles:

• loss_activity
  • BCEWithLogits on slot activity
  • computed over valid time steps
• loss_azi
  • cross-entropy over 360 azimuth bins
• loss_ele
  • cross-entropy over 180 elevation bins
• loss_dist
  • SmoothL1Loss on continuous distance regression
• loss_cls_aux
  • auxiliary source class cross-entropy
• loss_temp
  • temporal smoothness regularization over valid consecutive steps

The total loss is the weighted sum defined in SpatialLossConfig.

11. Stage-1 Training Flow

The current training entry is stage-1 encoder-focused training.

High-level flow per step:

batch
  -> SpatialBEATs.forward()
  -> match_fixed_slots()
  -> compute_spatial_losses()
  -> backward()
  -> optimizer.step()

11.1 Trainable modules in stage 1

By default, stage 1 trains:

• preprocessor
• patch_embedding
• frequency_pool
• temporal_resampler
• temporal_readout
• slot_readout
• prediction_heads

It can also unfreeze the BEATs trunk.

The projector is kept frozen by default in stage 1.

11.2 Optimizer

The current trainer uses:

• AdamW

Default preset values:

• batch_size = 4
• num_epochs = 20
• learning_rate = 1e-4
• weight_decay = 0.05

12. Current Presets

12.1 OV123_STAGE1_CFG

Defined in:

• spatial_beats_ov123_stage1_config.py

This preset is intended to train on:

• ov1_foa.jsonl
• ov2_foa.jsonl
• ov3_foa.jsonl

with split filtering:

• train: ("train",)
• val: ("valid",)
• test: ("test",)

and clip truncation:

• max_clip_duration_seconds = 20.0
• crop_mode = "start"

12.2 OV23_STAGE1_CFG

This is the safer baseline preset using only:

• ov2_foa.jsonl
• ov3_foa.jsonl

It uses the same split and truncation policy.

12.3 Important note on ov1

The trainer is already written to use split filtering for ov1.

If the active ov1 manifest at the configured path does not yet contain split, then:

• the OV123 preset will not automatically include those samples in train, valid, or test
• the fix is simply to point the preset at the updated ov1 manifest path

The code path itself already supports split-aware loading.

13. Current Runtime Status

The current implementation has already been checked on:

• a real FOA waveform file
• mixed-length real manifest samples
• full forward pass
• fixed-slot matching
• multi-task loss computation
• BEATs pretrained weight loading

The following paths are already operational:

• dataset parsing
• waveform batching
• mixed-length temporal masking
• model forward
• matching
• loss computation
• stage-1 optimization loop

14. Recommended Launch Pattern

Example usage:

from spatial_beats_ov123_stage1_config import OV123_STAGE1_CFG
from train_spatial_beats import main

main(OV123_STAGE1_CFG)

If ov1 still needs a different manifest path, update only:

OV123_STAGE1_CFG.train_manifest_paths
OV123_STAGE1_CFG.val_manifest_paths
OV123_STAGE1_CFG.test_manifest_paths

or rebuild the config through:

from train_spatial_beats import make_ov123_stage1_config

15. Summary

The current Spatial-BEATs implementation is a FOA-first BEATs-based spatial encoder with:

• Qwen-like low-level mel settings
• a 7-channel FOA front-end
• reused BEATs trunk
• fixed-rate 2.5 Hz spatial token output
• fixed-slot supervision heads
• variable-length batching
• split-aware ov1/ov2/ov3 training presets

The central training idea is:

• use explicit spatial supervision to shape the front-end and BEATs trunk
• keep the supervision head lightweight
• use encoder-side spatial embeddings as the final source of LLM spatial tokens