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.
| Attention type | Seed | Parameters | Eval loss | Checkpoint |
|---|---|---|---|---|
| MHA | 44 | 109,514,298 | 2.0972 | checkpoint-100000 |
| GQA | 42 | 100,064,826 | 2.1450 | checkpoint-100000 |
| MQA | 42 | 96,521,274 | 2.2177 | checkpoint-100000 |
| Collaborative MHA | 43 | 96,530,490 | 2.1196 | checkpoint-100000 |
| GT-MHA exact | 42 | 96,128,826 | 2.1428 | checkpoint-100000 |
| GT-MHA residual | 44 | 96,128,826 | 2.0961 | checkpoint-100000 |
The full 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.
Fine-tuning
Task-specific fine-tuning checkpoints for MNLI, QNLI, SST-2, and CoLA will be added in separate directories later. The current BERT-Base Checkpoints/ tree contains only the selected pretrained BERT-Base checkpoints.
Loading a checkpoint
from transformers import AutoModel, 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 = AutoModel.from_pretrained(repo, subfolder=path)
The repository is intended for research comparison and reproducibility. See the main code repository for model definitions, training scripts, and evaluation procedures.