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](/apdcephfs_cq10/share_1603164/user/schmittzhu/code/unilm/beats/spatial_beats.py)
- [spatial_modules.py](/apdcephfs_cq10/share_1603164/user/schmittzhu/code/unilm/beats/spatial_modules.py)
- [spatial_dataset.py](/apdcephfs_cq10/share_1603164/user/schmittzhu/code/unilm/beats/spatial_dataset.py)
- [spatial_loss.py](/apdcephfs_cq10/share_1603164/user/schmittzhu/code/unilm/beats/spatial_loss.py)
- [train_spatial_beats.py](/apdcephfs_cq10/share_1603164/user/schmittzhu/code/unilm/beats/train_spatial_beats.py)
- [spatial_beats_ov123_stage1_config.py](/apdcephfs_cq10/share_1603164/user/schmittzhu/code/unilm/beats/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:

```text
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:

```text
[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:

```text
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:

```text
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:

```text
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:

```text
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:

```text
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](/apdcephfs_cq10/share_1603164/user/schmittzhu/code/unilm/beats/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:

```python
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:

```python
OV123_STAGE1_CFG.train_manifest_paths
OV123_STAGE1_CFG.val_manifest_paths
OV123_STAGE1_CFG.test_manifest_paths
```

or rebuild the config through:

```python
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