Skip to content

Improve localization: ratio-based top_k, BF16 fix, and rename localize_fed10#402

Merged
KartikP merged 7 commits into
brain-score:mainfrom
BKHMSI:main
Jun 2, 2026
Merged

Improve localization: ratio-based top_k, BF16 fix, and rename localize_fed10#402
KartikP merged 7 commits into
brain-score:mainfrom
BKHMSI:main

Conversation

@BKHMSI
Copy link
Copy Markdown
Contributor

@BKHMSI BKHMSI commented May 28, 2026

  • Ratio support for top_k: localize_fedorenko2010 now accepts a float in (0, 1) as a fraction of total units (e.g. 0.1 = top 10%), in addition to absolute integers.
  • BF16 compatibility: Cast tensors to float32 before calling .numpy() in both _tensor_to_numpy and extract_representations, fixing a crash when running BF16 models.
  • Rename localize_fed10localize_fedorenko2010: More descriptive and consistent name; updated all call sites.
  • Refactor example script: score_localization.py is now a reusable function that auto-reads num_layers and hidden_size from the model config; updated default model to Qwen/Qwen3-4B.

Test plan

  • Run localization with a BF16 model (e.g. Qwen/Qwen3-4B) and confirm no dtype errors
  • Verify top_k=0.1 selects 10% of units and top_k=4096 still works as an absolute count
  • Confirm cached masks load correctly with the new model_id_path (slash-replaced) save path

@KartikP KartikP added OOM and removed OOM labels May 28, 2026
@KartikP KartikP closed this Jun 2, 2026
@KartikP KartikP reopened this Jun 2, 2026
@KartikP KartikP merged commit 857c9c1 into brain-score:main Jun 2, 2026
11 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants