Unmute the Patch Tokens

This is the official GitHub repository for the paper: Unmute the Patch Tokens: Rethinking Probing in Multi-Label Audio Classification, accepted at ICLR 2026.

🚀 Quick Start: Probing Methods

If you want to directly use the probing methods, not the complete pipeline.

Because the naming in the codebase differs slightly from the simplified names used in the paper, please refer to the mapping below:

Here is an example of how the input is extracted from the EAT model, reshaped, and processed before being fed into the probing head. In this implementation, the probes utilize the full patch tokens and subtract the class (cls) features to compute the relative features ($z_f$).

# extract features from EAT (normalization is handled internally)
features = model.extract_features(images) 
cls_token = features[:, 0]
patch_tokens = features[:, 1:]

# calculate dimensions and reshape patch tokens
B, _, T, F = images.shape
ps = model.model.local_encoder.proj.kernel_size[0]
H_ps, W_ps = T // ps, F // ps 

# shape becomes: (Batch, Hidden_Dim, H_ps, W_ps) e.g., (B, dim, 64, 8)
patch_features = patch_tokens.reshape(B, H_ps, W_ps, -1).permute(0, 3, 1, 2).contiguous() 

# compute relative patch features (z_f = x_patch - x_cls)
B, C, H, W = patch_features.shape
x_patch = patch_features.permute(0, 2, 3, 1).reshape(B, H * W, C)
z_f = x_patch - cls_token.unsqueeze(1)
probe_input_features = z_f.permute(0, 2, 1).reshape(B, C, H, W)

Setup

conda create -n probing python=3.11.9 -y
pip install -r requirements.txt

Overview

The workflow consists of three main steps:

1. Embedding computation: Extract and cache features from backbone models. While our caching method allows this to be done on-the-fly, we found it more practical to separate this step when running on Slurm.
2. Hyperparameter optimization (HPO): Use Optuna to find best hyperparameters.
3. Final training and results: Train models with optimized hyperparameters.

Note: Before running the workflow, make sure to adjust the data_path in the paths/cluster.yaml or paths/workstation.yaml file to the path where you want to store the data etc.

Stage 1: Embedding Computation

Pre-extract features from backbone models to speed up hyperparameter optimization by avoiding repeated feature extraction during HPO trials.

Usage

Single Configuration

python precompute_features.py \
  paths=cluster \
  backbone=audio/eat_base \
  dataset=audio/urban_sed

Batch Processing with SLURM

We used slurm to precompute features for all backbones on all datasets. See the configs/slurm/features directory for the slurm scripts. An example is given below.

sbatch configs/slurm/features/urban_sed.sh

Note: It could be easier to save the dataset to disk and then load it from disk as a separate step. See the configs/dataset/audio/urban_sed.yaml file for the configuration.

Configuration

Features are precomputed using the precompute.yaml config:

• Input: Audio from datasets
• Output: Cached features saved as .pth files in {data_path}/precache/{dataset_name}/{backbone_name}/
• Files created: train.pth, test.pth, valid.pth (if validation split exists)

Stage 2: Hyperparameter optimization (HPO)

Use Optuna with hybrid sampling (Sobol + TPE) to find optimal learning rates and weight decay values for different probing methods.

Usage

Single HPO Run

python optimize_hparams.py \
  optimization=twoway_sobol_prot \
  optimization.reset_study=True \
  trainer.logger=null \
  paths=workstation \
  head=prototypical \
  dataset=audio/urban_sed \
  backbone=audio/eat_base

Batch HPO with SLURM

Example:

sbatch configs/slurm/hpo/urban_sed/eat_base.sh

Configuration Options

Optimization Configs

• twoway_sobol: Standard HPO for most probing methods
• twoway_sobol_prot: Different search space for prototypical methods

Key Parameters

• n_trials: 50: Number of optimization trials
• exploration_ratio: 0.50: Fraction of trials using Sobol sampling
• reeval_topk: 3: Number of top configurations to re-evaluate
• reeval_seeds: [0,1,2,3,4]: Seeds for final evaluation

Search Space

search_space:
  optimizer:
    lr: [2e-3, 8e-2]           # Learning rate range (log scale)
    weight_decay: [1e-5, 5e-4] # Weight decay range (log scale)

Output

• Study results: Optuna trials saved to journal files or database
• Top-K parameters: JSON files with best hyperparameter configurations
• Re-evaluation results: Performance statistics across multiple seeds

Stage 3: Final training and results

Train models using the optimized hyperparameters across multiple random seeds for robust performance evaluation.

Usage

Extract Optimal Hyperparameters

From the re-evaluation JSON file, identify the best performing configuration. For instance:

{
  "rank": 1,
  "params": {
    "optimizer.lr": 0.0,
    "optimizer.weight_decay": 0.0
  },
  "mean": 0.0,
  "std": 0.0
}

Run final training

python main.py \
  paths=workstation \
  head=prototypical \
  dataset=audio/urban_sed \
  backbone=audio/eat_base \
  optimizer.lr=0.00 \
  optimizer.weight_decay=0.0 \
  seed=0

Multiple Seeds

# Run across different seeds for statistical significance
for seed in 0 1 2 3 4; do
  python main.py \
    paths=workstation \
    head=prototypical \
    dataset=audio/eat_base \
    backbone=audio/urban_sed \
    optimizer.lr=0.0 \
    optimizer.weight_decay=0.0 \
    seed=$seed \
    experiment_name="final_eval_${seed}"
done

Complete Workflow Example

Here's a complete example for the urban_sed dataset with EAT backbone:

# 1. Precompute features
python precompute_features.py \
  paths=workstation \
  backbone=audio/eat_base \
  dataset=audio/urban_sed

# 2. Run hyperparameter optimization
python optimize_hparams.py \
  optimization=twoway_sobol_prot \
  optimization.reset_study=True \
  trainer.logger=null \
  paths=workstation \
  head=prototypical \
  dataset=audio/urban_sed \
  backbone=audio/eat_base

# 3. Extract best hyperparameters from JSON output
# (manually inspect reeval_top3_*.json file)

# 4. Run final training with optimal hyperparameters
python main.py \
  paths=workstation \
  head=prototypical \
  dataset=audio/urban_sed \
  backbone=audio/eat_base \
  optimizer.lr=0.00 \
  optimizer.weight_decay=0.0 \
  seed=0