GT-MHA

Group-Transformed Multi-Head Attention (GT-MHA)

This repository contains pretrained BERT-Base checkpoints for comparing attention mechanisms, with Group-Transformed Multi-Head Attention (GT-MHA) as the main contribution of this work.

GT-MHA transforms groups of attention heads while retaining the multi-head structure. The repository includes both the exact formulation and a residual formulation, alongside standard attention baselines for controlled comparisons.

BERT-Base checkpoints

The table reports the lowest recorded validation (eval_loss) value during masked-language-model pretraining. Each entry points to the corresponding saved checkpoint directory.

The selected pretrained checkpoint trees are organized under BERT-Base Checkpoints/:

BERT-Base Checkpoints/
├── mha/
├── gqa/
├── mqa/
├── collaborative_mha/
├── gt_mha_exact/
└── gt_mha_residual/

Each selected checkpoint includes the model weights, configuration, tokenizer, trainer state, optimizer and scheduler state, RNG state, and training arguments. This makes the runs usable both for evaluation and for exact training resumption.

Official GLUE test results

The table reports hidden-test results returned by the official GLUE evaluation server on August 29, 2026. Each task model was independently fine-tuned from its architecture's pretrained checkpoint. Hyperparameters and checkpoints were selected using validation data only; official test labels were never accessed or used for model selection. Scores are on a 0–100 scale. Paired metrics are reported in the same order as the official server.

Under the GLUE submission convention, which excludes word or word-part embeddings, MHA has 86,043,651 parameters, GQA has 76,594,179, collaborativeMHA has 73,059,843, and GT-MHA residual has 72,658,179. GQA and collaborativeMHA use four key-value heads and shared query-key projections, respectively, while all task models are fine-tuned independently.

Fine-tuning and checkpoint-selection protocol

The eight primary tasks used an equal-budget task-specific learning-rate sweep for every evaluated attention type, with effective batch size 128 and learning rates {2e-5, 3e-5, 5e-5}. MNLI, QNLI, SST-2, and QQP used four epochs; CoLA, MRPC, STS-B, and RTE used eight epochs. For each architecture and task, the learning rate was selected by its mean primary validation metric over seeds 42–44, followed by confirmation runs with seeds 45 and 46. The submitted checkpoint was the highest-scoring seed at that preselected learning rate. Validation-metric ties were broken by lower validation loss.

WNLI was not swept because its tiny, adversarial validation split makes architecture-specific hyperparameter selection unreliable. Its submission checkpoint used the fixed common setting of learning rate 2e-5, seed 42, and four epochs.

TinyStories-v2 checkpoints

The table reports validation loss at step 250000 (the selected checkpoint may be a nearby saved step where noted).

TinyStories checkpoints are organized under TinyStories-v2 Checkpoints/.

Loading a checkpoint

from lgma.bert import load_bert_sequence_classifier
from transformers import AutoTokenizer

repo = "shahils/GT-MHA"
path = "BERT-Base Checkpoints/gt_mha_exact/bert_base_gt_mha_exact_b4g8h12_random_seed42_explicit_fuseqkv_v2/checkpoint-100000"

tokenizer = AutoTokenizer.from_pretrained(repo, subfolder=path)
model, audit = load_bert_sequence_classifier(
    "/local/path/to/checkpoint-100000",
    num_labels=2,
    attention_type="gt_mha_residual",
)

For non-MHA checkpoints, keep bert_gt_mha_manifest.json and gt_mha_state_dict.pt in the same checkpoint directory as config.json and model.safetensors. The loader uses these files to restore the trained attention modules; loading only model.safetensors with a generic Transformers class will fall back to newly initialized attention.

The repository is intended for research comparison and reproducibility. See the main code repository for model definitions, training scripts, and evaluation procedures.