[Proposal] Add Audio Spectrogram Transformer (AST) Adapter (ASTForAudioClassification) #1446

Proposal

Add a TransformerBridge adapter for the Audio Spectrogram Transformer family (ASTForAudioClassification)

Motivation

TransformerLens currently supports audio through the HookedAudioEncoder (HuBERT/Wav2Vec2), and this proposal would add an adapter for the Vision-Transformer (ViT) based audio models. The AST is widely used and processes spectrograms, which are 2D images representing audio data, making it a prime candidate for Mechanistic Interpretability & TransformerLens, expanding the library's multi modal capabilities.

Pitch

I plan to implement the AST adapter for the TransformerBridge architecture, including:

• mapping the HuggingFace AST configuration to HookedTransformerConfig
• writing the weight translation map for the dense ViT-style attention and MLP blocks
• implementing the patch embedding extraction layer (handling the 2D Conv2D projection)to ensure alignment between input spectrograms and the residual stream before the transformer blocks

I have a fair amount of experience with this particular model, and will gladly work on the implementation and testing for this adapter.

Alternatives

I've considered doing this for ViT-base given its backbone is functionally identical to the AST (the AST input layers are altered, but the models' innards are the same), but decided to start with the AST since I'd be surprised if we don't have this for ViT-base already. Adapting the work of this proposal for ViT-base in the future should be simple.

Additional context

HF for AST model: https://huggingface.co/docs/transformers/model_doc/audio-spectrogram-transformer

Checklist

• I have checked that there is no similar issue in the repo (required)