Add AST (Audio Spectrogram Transformer) Adapter

[dylanberens]

Description

Refactored the AST architecture adapter to align with the V3 'TransformerBridge' architecture and target 'dev'.

Key Changes

• V3 Bridge Architecture: Implemented 'ASTArchitectureAdapter' using hierarchical 'component_mapping' ('BlockBridge', 'AttentionBridge', 'MLPBridge', 'NormalizationBridge', and 'UnembeddingBridge') & added ASTForAudioClassification to architecture adapter factory
• Weight Processing: Ported head-splitting tensor reshapes ('RearrangeTensorConversion') for Query/Key/Value/Output projections to handle dimensional conversion during loading
• Dynamic Context Length: Utilized 'prepare_model' hook to dynamically extract 'n_ctx' from positional embeddings length
• Unit Test: Updated unit tests under 'tests/unit/model_bridge/supported_architectures/test_ast_adapter.py' asserting exact bridge-vs-HF parity

Fixes # 1484

Type of change

• New feature (non-breaking change which adds functionality)

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

[jlarson4]

Thanks for putting this together @dylanberens! The AST domain work is great, the Q/K/V/O and bias reshapes, MLP transposes, the overlapping-patch Conv2d, the cls-distillation-patches token order, and the pooled dual-token + dual-LayerNorm head are all correct.

We will need to address & adjust the approach: this is written against the legacy HookedTransformer convert_weights path, but TransformerBridge adapters are driven by component_mapping (wrap HF's live modules). As submitted it can't load or run, which is why your tests are failing in CI.

Primary issues with the current code:

1. Doesn't integrate with the bridge – ASTAdapter never sets self.component_mapping (it defines no __init__, so it inherits None), and it isn't registered in the factory. Booting AST raises ValueError: Unsupported architecture. get_config_map/convert_weights are legacy-HT classmethods no bridge path ever calls.
2. The test fails and never touches the bridge – It KeyErrors at ast.py:125, the HF keys are wrong on two axes: layers.{i} should be encoder.layer.{i}, and q_proj/k_proj/v_proj/o_proj/mlp.fc1/fc2 should be attention.attention.{query,key,value} / attention.output.dense / intermediate.dense / output.dense. So the max diff < 1e-4 claim isn't reproducible, and the test builds a raw HookedTransformer, not a bridge.

To make this mergeable:

• Rewrite as a component_mapping adapter: I'd recommend following our Architecture Adapter Creation Guide. The hubert.py adapter is our current working audio adapter, it would be a good reference point to work off.
• Delete ASTEmbed and both classmethods — the bridge wraps HF's own modules, so none of that math needs reimplementing.
• Fix n_ctx = num_patches+2; read patch/stride/hidden_act from config.
• Register ASTForAudioClassification in the factory + export it from __init__.
• Rewrite the test to boot a real TransformerBridge and assert bridge-vs-HF fp32 parity, under tests/unit/....

Also, a couple repo items: make sure to drop poetry.lock from the PR. Starting from 3.0.0, the repo is uv-based. And, base this change to dev, we will be able to launch this as part of a 3.x release, it does not need to be gated behind version 4.

Happy to re-review once it's on the TransformerBridge pattern.

[dylanberens]

thanks @jlarson4 for the feedback and roadmap! I've refactored the AST adapter to follow the V3 TransformerBridge pattern against dev

• removed custom embedding reimplementations and adopted hierarchical component_mapping
• ported q/k/v/o head splitting tensor conversions
• utilized prepare_model to resolve n_ctx from positional embeddings
• rebase to dev, added the ASTForAudioClassification to architecture adapter factory, added FP32 bridge parity unit tests (passing locally with exact 0.0 difference)

ready for re-review when you get a chance, thanks jonah!

[jlarson4]

Excellent, thank you @dylanberens! I will make sure to let you know when I get around to reviewing

[jlarson4]

Hey @dylanberens sorry it took me so long to get around to reviewing this. Full detailed review broken down below. The overall structure and content of your adapter are very good! There are just a couple small items that need addressing before I can merge. Once the comments below are resolved, this should be good to go!

[jlarson4]

This is the only adapter class in supported_architectures/) without a docstring, and the file also lacks a module docstring. AST has enough quirks that it would really benefit from one: the spectrogram input contract, the cls+distillation token layout, and especially the head structure.

It also needs the encoder-only class attribute:

class ASTArchitectureAdapter(ArchitectureAdapter):
    """Architecture adapter for AST (Audio Spectrogram Transformer) audio classifiers.

    Input is a [batch, time=1024, n_mels=128] spectrogram; positions 0/1 are the
    CLS and distillation tokens. HF pools (cls+dist)/2 after ln_final and applies
    an extra LayerNorm inside the classifier head, see the unembed mapping note.
    """

    supports_generation: bool = False

Without it, AST inherits True and every generation entry point will fail with unclear errors instead of the clean NotImplementedError gate at bridge.py:2721. I tested this and received AssertionError: Must pass eos_token_id... from generate(spectrogram), and IndexError at bridge.py:2631 with stop_at_eos=False.

[jlarson4]

HuggingFace's classification path after ln_final is: pool (seq[:,0]+seq[:,1])/2 -> classifier.layernorm -> classifier.dense. Mapping unembed to classifier.dense alone means:

• unembed.hook_in is the pooled, post-head-LayerNorm [batch, d_model] tensor. There are two unhooked ops downstream of ln_final.hook_out [batch, seq, d_model]. On every text model, unembed input is the ln_final output. Users expecting that convention get silently different semantics here.
• The pooling step and classifier.layernorm have no hook points at all, neither is observable or intervenable.

The actual results are unaffected, so I'm fine shipping this mapping if documented in the class docstring, but consider mapping the head LN as a component (or the whole classifier) so the head is hookable.

Either way, a couple related small bugs:

• ASTConfig(num_labels=0) makes classifier.dense an nn.Identity, and boot then dies with AssertionError: Component classifier.dense has no weight attribute (bridge.py:667). Guard it in prepare_model (drop the unembed mapping, like HuBERT's conditional unembed injection) or raise a clear unsupported-config error.
• cfg.d_vocab/d_vocab_out stay at the -1 sentinel after boot even though W_U is a real (d_model, num_labels) head. Set self.cfg.d_vocab_out (and arguably d_vocab) from num_labels so downstream consumers don't read -1.

[jlarson4]

The dynamic n_ctx is the right call, but the hardcoded hf_model.audio_spectrogram_transformer. access is one of the two things that breaks real boots: a bare ASTModel has no such attribute.

Please make this and the component_mapping prefix-aware following hubert.py's _build_component_mapping(prefix=...) pattern. It solves the bare-encoder case and a real boot blocker. This is also the natural place for the num_labels -> d_vocab_out propagation mentioned above.

[jlarson4]

We could use a bit more coverage from the testing. We have a guide on what we typically expect from adapter tests here. Let me know if you have any questions!

[jlarson4]

Because the adapter sets is_audio_model=True, boot reaches the processor block (sources/transformers.py:977-989) and calls AutoFeatureExtractor.from_pretrained("ast"). This attempts to find a HF repo literally named ast on every unit-test run. If anyone ever publishes a repo named ast with a preprocessor_config.json, this test silently loads third-party content. Use an unmistakably-local placeholder and/or monkeypatch the processor load.

[jlarson4]

Drop the print()s (pytest reporting covers it) and the `if name == "main": block. The TransformerLens standard are assertions with messages.

[jlarson4]

There are a couple other locations you'll need to register this:

1. AUDIO_ARCHITECTURES in transformer_lens/utilities/architectures.py needs "ASTForAudioClassification". This is routes verify_models to the audio phase set {1, 8} instead of text phases (see review body) when running benchmarks to confirm individual model support.
2. ARCHITECTURE_DESCRIPTIONS in generate_report.py this is a nice-to-have, not essential for PR merge