[Proposal] RoBERTa masked-LM adapter for TransformerBridge #1870

[Proposal] RoBERTa masked-LM adapter for TransformerBridge

Proposal

Add a RobertaArchitectureAdapter so RobertaForMaskedLM checkpoints (roberta-base, roberta-large, distilroberta-base and the many fine-tunes that keep the MLM head) load through TransformerBridge.boot_transformers.

Right now this fails at adapter selection:

ValueError: Unsupported architecture: RobertaForMaskedLM

even though transformer_lens/utilities/architectures.py already lists RobertaForMaskedLM under MASKED_LM_ARCHITECTURES and routes it to the masked-LM loader. The classifier knows about the family; but the factory doesn't.

Motivation

BERT is the only encoder-only text model the bridge supports. For anything that compares pretraining recipes, or reuses an existing probing / patching setup on a model that isn't BERT, RoBERTa is usually the next model people reach for, and it's currently a dead end in both HookedEncoder and the bridge. The concrete case is running the same activation-patching or probing code on bert-base-cased, roberta-base and distilroberta-base without changing anything but the model id.

I searched issues and PRs for "roberta" (open and closed) and found nothing. #918 added DistilBERT to the legacy loading_from_pretrained path and was closed unmerged in May 2025; it never touched the bridge, so I don't think it counts as prior art here.

Pitch

This is close to a prefix rename of the existing BERT adapter. I instantiated a tiny RobertaForMaskedLM and diffed its module tree against what bert.py maps:

Everything inside a layer (attention.self.{query,key,value}, attention.output.dense, both LayerNorms, intermediate.dense, output.dense) is identical, so the per-block submodules, the q/k/v/o rearrange conversions, supports_fold_ln = False for post-LN and supports_generation = False all carry over as-is. The decoder weight is tied to the word embeddings in both models.

So the plan would be:

1. supported_architectures/roberta.py with RobertaArchitectureAdapter(BertArchitectureAdapter) that only overrides component_mapping with the new names. Same shape as Olmo3ArchitectureAdapter(Olmo2ArchitectureAdapter).
2. The four registration sites from supported_architectures/AGENTS.md: __init__.py, SUPPORTED_ARCHITECTURES in the factory, HF_SUPPORTED_ARCHITECTURES + CANONICAL_AUTHORS_BY_ARCH (["FacebookAI", "distilbert"]) in tools/model_registry/__init__.py, and ARCHITECTURE_DESCRIPTIONS in generate_report.py (that one already has a RoBERTa string, it just isn't reachable).
3. supported_models.json entries for FacebookAI/roberta-base and distilbert/distilroberta-base with status: 0, then a verify_models run on both. distilroberta is 6 layers / 82M params so it's cheap to iterate on; roberta-base is the one people will actually use.
4. Tests: a unit adapter test next to test_bert_adapter.py, and an integration test asserting HF logit parity, hook caching and a basic intervention on distilroberta. There is no bridge-side BERT integration test to copy (the BERT ones under tests/integration and tests/acceptance are HookedEncoder tests), so this would be the first one for an encoder.

Two things that are not a rename and need a decision:

• n_ctx. sources/transformers.py sets n_ctx = max_position_embeddings, which is 514 for RoBERTa because HF reserves positions 0 and 1 (padding_idx + 1 offset). Usable context is 512. The adapter should correct this so n_ctx matches what you can actually feed the model. Related: hook_pos_embed will see position ids starting at 2, not 0, since HF computes them inside RobertaEmbeddings. The bridge runs the HF forward so the values are right, but anyone comparing position-embedding activations against BERT should know about the offset; Worth a line in the adapter docstring.
• Tokenizer policy. RoBERTa is byte-level BPE with <s> / </s> rather than WordPiece with [CLS] / [SEP]. Per the per-model tokenizer rule in supported_architectures/AGENTS.md, default_prepend_bos and padding side should come from each checkpoint's tokenizer_config.json, not be inherited from the BERT adapter.

Alternatives

• Load through transformers directly and register forward hooks by hand. Works, but it's exactly the thing the bridge exists to avoid, and you lose run_with_cache, the hook aliases, and compatibility mode.
• Port RoBERTa to HookedEncoder. AGENTS.md says the bridge is the default for new work, and the BERT NSP note in migrating_to_v3.md already points people at the bridge for new encoder features, so this seems like the wrong direction.
• A single generic "BERT-like" adapter keyed on prefix., which is tempting, but the head module names differ enough between BERT, RoBERTa, ALBERT and ELECTRA that I think one subclass per family is clearer and matches how the repo handles other sibling families.

Additional context

• distilbert/distilroberta-base and FacebookAI/roberta-base both declare architectures: ["RobertaForMaskedLM"], model_type: roberta, pad_token_id: 1, type_vocab_size: 1, max_position_embeddings: 514. They differ only in num_hidden_layers (6 vs 12).
• token_type_embeddings exists in both BERT and RoBERTa but with type_vocab_size = 1 in RoBERTa it's a constant. The BERT adapter doesn't map it either, so no change there.
• Happy to do the PR, based on dev, following the adapter workflow in supported_architectures/AGENTS.md. I can run verification for both models locally. Mostly opening this to check nobody else has it in flight and that a BertArchitectureAdapter subclass is the shape you want.

Checklist

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