brain-score / brain-score/vision

Layer accumulation in the activation cache is incompatible with transformer layer sets

Open
#2,514 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
193
Forks
105
Avg merge
10h 48m
Merged PRs (30d)
10

Description

The stored-activation cache accumulates layers into a single entry via `combine_fields={'layers': 'layer'}`. When a later call requests a layer that isn't stored yet, `_XarrayStorage` computes it and merges it into the existing entry with `xr.merge`.

That merge fails when two layers produce different neuroid coordinate schemes. `core.py:311-317` names coords by flatten rank:

| rank | example | coords |
|---|---|---|
| 1 | classifier / FC | `channel`, `channel_x`, `channel_y` |
| 2 | **transformer block** | `channel`, `embedding` |
| 3 | 2D conv | `channel`, `channel_x`, `channel_y` |

Ranks 1 and 3 agree, which is why CNNs are unaffected. Rank 2 does not, so a transformer scored on both a neural benchmark (mapped to a block) and a behavioural one (mapped to the classifier) tries to merge `embedding` with `channel_x`/`channel_y`.

Reproduced on current master, xarray 2022.3.0:

```
AssertionError: Length of new_levels (5) must be <= self.nlevels (4)
```

## Why this has been invisible

Two plugin families work around it by disabling the activation cache for ViT configs — `scaling_models` and `blur_timm_models`, both citing #1232, which was **closed, not merged**. So the "temporary fix" has been the only thing holding since, is undocumented outside those two files, and every future transformer submission has to rediscover it independently.

It also didn't surface in production until recently: scoring ran with `RESULTCACHING_DISABLE=1`, so layer merges never happened at all. With the cache now on by default, the merge path is live — though no failures are attributable to it yet (1,006 score rows since Aug 12; the single merge-shaped failure in the DB is `VOneCORnet-S` from January 2025, unrelated).

## Suggested direction

Options, roughly increasing in invasiveness:

1. **Refuse to accumulate across incompatible schemes** — detect the coord-name mismatch before merging and store the layer under a separate entry, rather than raising. Keeps caching for transformers, loses cross-layer accumulation for them.
2. **Normalise the coordinate names** so all ranks share a scheme. Cleaner, but changes assembly coords that downstream metrics may read.
3. **Keep disabling per-model**, but make it discoverable — a documented helper rather than an env-var mutation copied between plugins.

(3) is now partially done: `disable_activation_caching` in `model_helpers/activations/pytorch.py` scopes the disable to a model's own calls instead of setting a process-global that leaked to every model loaded afterwards in the same session.

## Related

- #1232 — the original ViT merge report (closed unmerged)
- #2503 / #2511 — preprocessing resolution in the cache key
- #2513 — wrapper identifiers unconstrained

Contributor guide

No contributing guide indexed for this repository

Research direction

Start with core.py:311-317 and the _XarrayStorage merge path, then read disable_activation_caching in model_helpers/activations/pytorch.py and the workarounds in scaling_models and blur_timm_models. Choose and document a consistent handling strategy for incompatible coordinate schemes, and verify that cached transformer and classifier layers no longer fail during accumulation.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
backend, data
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Needs clarification
Newbie friendliness
42/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.