Skip to content

Add AST (Audio Spectrogram Transformer) Adapter #1484

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
dylanberens wants to merge 4 commits into
base: dev
Choose a base branch
from
Open

Add AST (Audio Spectrogram Transformer) Adapter #1484

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
43d9ff4
chore: save WIP on V3 transformerbridge migration
dylanberens Jul 11, 2026
3b9db10
feat(ast): migrate AST adapter to V3 TransformerBridge and component_…
dylanberens Jul 25, 2026
2820b13
chore(ast): register ASTForAudioClassification and MIT canonical auth…
dylanberens Jul 25, 2026
3a2a164
Merge branch 'dev' into add-ast-adapter
jlarson4 Jul 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 56 additions & 0 deletions tests/unit/model_bridge/supported_architectures/test_ast_adapter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,56 @@
import torch
from transformers import ASTConfig, ASTForAudioClassification

from transformer_lens.model_bridge import TransformerBridge


def test_ast_parity():

@jlarson4 jlarson4 Jul 27, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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!

# 1. setup a tiny HF AST config for instantaneous local testing
hf_config = ASTConfig(
hidden_size=12,
num_hidden_layers=2,
num_attention_heads=2,
intermediate_size=24,
patch_size=16,
)

# explicitly define the architecture list for the local config
hf_config.architectures = ["ASTForAudioClassification"]

hf_model = ASTForAudioClassification(hf_config)
hf_model.eval()

# 2. boot the v3 transformerbridge natively
# bridge auotmatically detecty ASTForAudioClassification and use new ASTArchitectureAdapater
tl_model = TransformerBridge.boot_transformers("ast", hf_model=hf_model)

@jlarson4 jlarson4 Jul 27, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

tl_model.eval()

# 3. create dummy spectrogram input [batch, freq, time]
dummy_spectrogram = torch.randn(1, 1024, 128)

# 4. compare forward passes
with torch.no_grad():
hf_logits = hf_model(dummy_spectrogram).logits

# bridge forward pass (extract sequence before HF's pooling)
_, cache = tl_model.run_with_cache(dummy_spectrogram)
resid = cache["ln_final.hook_normalized"]

# apply AST pooling logic to the bridges residual stream
cls_token = resid[:, 0, :]
dist_token = resid[:, 1, :]
pooled_out = (cls_token + dist_token) / 2.0

# apply the final classifier (which holds 2nd layernorm + dense)
tl_logits = hf_model.classifier(pooled_out)

diff = (hf_logits - tl_logits).abs().max().item()
print(f"Bridge vs HF Max Logit Diff: {diff:.6e}")

@jlarson4 jlarson4 Jul 27, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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


# 5. assert parity
assert diff < 1e-4, "Parity failed: Bridge tensors do not match HuggingFace."
print("Parity dub. V3 Bridge connected excellent")


if __name__ == "__main__":
test_ast_parity()
2 changes: 2 additions & 0 deletions transformer_lens/factories/architecture_adapter_factory.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
AfmoeArchitectureAdapter,
ApertusArchitectureAdapter,
ArceeArchitectureAdapter,
ASTArchitectureAdapter,
AudioFlamingo3ArchitectureAdapter,
BaichuanArchitectureAdapter,
BambaArchitectureAdapter,
Expand Down Expand Up @@ -158,6 +159,7 @@
"AfmoeForCausalLM": AfmoeArchitectureAdapter,
"ApertusForCausalLM": ApertusArchitectureAdapter,
"ArceeForCausalLM": ArceeArchitectureAdapter,
"ASTForAudioClassification": ASTArchitectureAdapter,
"BaiChuanForCausalLM": BaichuanArchitectureAdapter,
"BaichuanForCausalLM": BaichuanArchitectureAdapter,
"BambaForCausalLM": BambaArchitectureAdapter,
Expand Down
4 changes: 4 additions & 0 deletions transformer_lens/model_bridge/supported_architectures/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,9 @@
from transformer_lens.model_bridge.supported_architectures.audio_flamingo3 import (
AudioFlamingo3ArchitectureAdapter,
)
from transformer_lens.model_bridge.supported_architectures.ast import (
ASTArchitectureAdapter,
)
from transformer_lens.model_bridge.supported_architectures.baichuan import (
BaichuanArchitectureAdapter,
)
Expand Down Expand Up @@ -271,6 +274,7 @@
"AfmoeArchitectureAdapter",
"ApertusArchitectureAdapter",
"ArceeArchitectureAdapter",
"ASTArchitectureAdapter",
"AudioFlamingo3ArchitectureAdapter",
"BD3LMArchitectureAdapter",
"BaichuanArchitectureAdapter",
Expand Down
103 changes: 103 additions & 0 deletions transformer_lens/model_bridge/supported_architectures/ast.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,103 @@
from typing import Any

from transformer_lens.conversion_utils.conversion_steps import RearrangeTensorConversion
from transformer_lens.conversion_utils.param_processing_conversion import (
ParamProcessingConversion,
)
from transformer_lens.model_bridge.architecture_adapter import ArchitectureAdapter
from transformer_lens.model_bridge.generalized_components import (
AttentionBridge,
BlockBridge,
LinearBridge,
MLPBridge,
NormalizationBridge,
UnembeddingBridge,
)
from transformer_lens.model_bridge.generalized_components.base import (
GeneralizedComponent,
)


class ASTArchitectureAdapter(ArchitectureAdapter):

@jlarson4 jlarson4 Jul 27, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

def __init__(self, cfg: Any) -> None:
super().__init__(cfg)

# essential flag for audio models in V3
self.cfg.is_audio_model = True
self.cfg.normalization_type = "LN"

n_heads = self.cfg.n_heads

# Q/K/V/O rearrangement: splits hidden dims into (heads, head_dim)
self.weight_processing_conversions = {
"blocks.{i}.attn.q.weight": ParamProcessingConversion(
tensor_conversion=RearrangeTensorConversion(
"(h d_head) d_model -> h d_model d_head", h=n_heads
),
),
"blocks.{i}.attn.k.weight": ParamProcessingConversion(
tensor_conversion=RearrangeTensorConversion(
"(h d_head) d_model -> h d_model d_head", h=n_heads
),
),
"blocks.{i}.attn.v.weight": ParamProcessingConversion(
tensor_conversion=RearrangeTensorConversion(
"(h d_head) d_model -> h d_model d_head", h=n_heads
),
),
"blocks.{i}.attn.q.bias": ParamProcessingConversion(
tensor_conversion=RearrangeTensorConversion("(h d_head) -> h d_head", h=n_heads),
),
"blocks.{i}.attn.k.bias": ParamProcessingConversion(
tensor_conversion=RearrangeTensorConversion("(h d_head) -> h d_head", h=n_heads),
),
"blocks.{i}.attn.v.bias": ParamProcessingConversion(
tensor_conversion=RearrangeTensorConversion("(h d_head) -> h d_head", h=n_heads),
),
"blocks.{i}.attn.o.weight": ParamProcessingConversion(
tensor_conversion=RearrangeTensorConversion(
"d_model (h d_head) -> h d_head d_model", h=n_heads
),
),
}

# V3 bridge pattern: hierarchical mapping using bridge components
self.component_mapping = {
"embed": GeneralizedComponent(name="audio_spectrogram_transformer.embeddings"),
"ln_final": NormalizationBridge(
name="audio_spectrogram_transformer.layernorm", config=self.cfg
),
"unembed": UnembeddingBridge(name="classifier.dense"),

@jlarson4 jlarson4 Jul 27, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

"blocks": BlockBridge(
name="audio_spectrogram_transformer.layers",
submodules={
"ln1": NormalizationBridge(name="layernorm_before", config=self.cfg),
"ln2": NormalizationBridge(name="layernorm_after", config=self.cfg),
"attn": AttentionBridge(
name="attention",
config=self.cfg,
submodules={
"q": LinearBridge(name="q_proj"),
"k": LinearBridge(name="k_proj"),
"v": LinearBridge(name="v_proj"),
"o": LinearBridge(name="o_proj"),
},
),
"mlp": MLPBridge(
name="mlp",
config=self.cfg,
submodules={
"in": LinearBridge(name="fc1"),
"out": LinearBridge(name="fc2"),
},
),
},
),
}

def prepare_model(self, hf_model: Any) -> None:
# hook to access the live Huggingface model before boot
# calculate n_ctx dynamically from the instantiated position embeddings
self.cfg.n_ctx = (

@jlarson4 jlarson4 Jul 27, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

hf_model.audio_spectrogram_transformer.embeddings.position_embeddings.shape[1]
)
2 changes: 2 additions & 0 deletions transformer_lens/tools/model_registry/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,7 @@
"AfmoeForCausalLM",
"ApertusForCausalLM",
"ArceeForCausalLM",
"ASTForAudioClassification",

@jlarson4 jlarson4 Jul 27, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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

"BaiChuanForCausalLM",
"BaichuanForCausalLM",
"BambaForCausalLM",
Expand Down Expand Up @@ -200,6 +201,7 @@
"AfmoeForCausalLM": ["arcee-ai"],
"ApertusForCausalLM": ["swiss-ai"],
"ArceeForCausalLM": ["arcee-ai"],
"ASTForAudioClassification": ["MIT"],
"BaiChuanForCausalLM": ["baichuan-inc"],
"BaichuanForCausalLM": ["baichuan-inc"],
"BambaForCausalLM": ["ibm-ai-platform"],
Expand Down
Loading