Skip to content

Add AST (Audio Spectrogram Transformer) Adapter#1484

Open
dylanberens wants to merge 4 commits into
TransformerLensOrg:dev-4.xfrom
dylanberens:add-ast-adapter
Open

Add AST (Audio Spectrogram Transformer) Adapter#1484
dylanberens wants to merge 4 commits into
TransformerLensOrg:dev-4.xfrom
dylanberens:add-ast-adapter

Conversation

@dylanberens

@dylanberens dylanberens commented Jul 3, 2026

Copy link
Copy Markdown

Description

This PR introduces full support for the Audio Spectrogram Transformer (AST) architecture with TransformerLens.

Key Engineering Details
Because AST is built on a ViT-base backbone but processes continuous audio spectrogrmas, several important architectural bridges were necessary:

  • Custom ASTEmbed: Implemented a specialized Conv2d embedding layer to handle 1D to 2D tensor reshaping and extract overlapping patches using a stride=(10, 10), which is a frequent cause of disconnects whens working with AST attention weights
  • Dual-Token Pooling: Implemented the AST-specific MLP head pooling logic, which averages botht eh CLS and Distillation tokens before classification.
  • Dual LayerNorm: Accounted for AST's second LayerNorm applied exclusively to the pooled token prior to the embedding matrix

Adds ast.py (architecture adapter) and test_ast_adapter.py (verifying FP32 mathematical parity with max logit diff < 1e-4).

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

map HF ASTConfig to HookedTransformerConfig

build custom ASTEmbed w/ 16x16 Conv2d patch extraction

translate q/k/v_proj and fc1/fc2 dense weights to TL standard blocks

implement CLS and Distillation token averaging for AST MLP head

apply dual layernorm architecture to pooled tokens

draft local matrix parity test suite (current max diff: ~0.05)
@dylanberens dylanberens changed the title [WIP] Add AST Adapter Add AST (Audio Spectrogram Transformer) Adapter Jul 5, 2026
@dylanberens dylanberens marked this pull request as ready for review July 5, 2026 01:49
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.

1 participant