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]

Thanks @dylanberens for the update! The fixes look good, I verified them all locally. Three things left before merge:

1. The real load path still can't reach your fixes – AUDIO_ARCHITECTURES routes AST to bare AutoModel, so a real boot_transformers("MIT/ast-...") loads a headless ASTModel without the classifier, and the audio_spectrogram_transformer. prefix branch never fires. Visual Encoders (ViT, DeiT) Support Rollout #1546 added an perfect example: mirror VISION_CLASSIFICATION_ARCHITECTURES -> AutoModelForAudioClassification in utilities/architectures.py and sources/transformers.py. Please also add a boot test, boot_transformers(<checkpoint>, load_weights=False) is cheap and fails today, so it'll pin the fix.

2. "ASTModel" in AUDIO_ARCHITECTURES is currently inert – it isn't registered in the factory or registry, so bare-encoder checkpoints still fail at adapter selection while the docstring claims support. Either register it in the other sites or drop the entry and narrow the docstring.

3. The parity test belongs in the integration folder (tests/integration/model_bridge/test_ast_tiny.py). A new adapter needs one there regardless.

[dylanberens]

@jlarson4 most recent commit split parity to integration folder, add load_weights boot test, fix audio classification load path, and add to ARCHITECTURE_DESCRIPTIONS in generate_report.py just to be thorough

also for ASTModel i opted to drop the entry and narrow the docstring, thanks man

[jlarson4]

Thanks @dylanberens! Great work on this, boot_transformers is now loading and returning logits as expected! A few new requests below. I appreciate your hard work on this, sorry for all the back and forth