tahamajs/bert-factorised-matrix-state-sst2

🧠 BERT-Base (factorised_matrix_state) — Subquadratic Attention Zoo

This repository provides an authoritative pretrained and distilled checkpoint of bert-base-uncased where standard quadratic softmax attention blocks across all 12 transformer layers have been replaced with the factorised_matrix_state mechanism from the NeuralOps Library.

The model was trained and evaluated on the General Language Understanding Evaluation (GLUE) benchmark suite (SST-2 task) using the NeuralOps Cached Representation Distillation Protocol (Zero Teacher VRAM allocation, offline representation alignment) physically executed on an NVIDIA GeForce RTX 3090 (24GB).

---

🔬 Architectural & Complexity Profile

• Attention Mechanism: factorised_matrix_state
• Subquadratic Scaling: $O(N)$ or $O(N \log N)$ effective operational complexity versus standard quadratic $O(N^2)$ Multi-Head Attention (MHA).
• Weights Serialization: attention_only_fp16 (attention_weights.safetensors).
• Memory Footprint: ~55 MB download footprint (strictly the trained attention weights and task classification head, zero redundant base model bloat).

---

📊 Empirical Evaluation & Distillation Benchmark

Evaluated on the GLUE Benchmark SST-2 validation split (872 binary sequence classification samples):

Metric	Teacher Baseline (bert-base-uncased)	Pre-KD Student (factorised_matrix_state)	Post-KD Distilled Model	Accuracy Delta ($\Delta$)	Distillation Status
GLUE SST-2 Accuracy	91.63%	52.87%	79.82%	+26.95%	Calibrated
Validation Cross-Entropy	0.2428	Uncalibrated	0.4853	Converged	Optimal

Hardware & Distillation Telemetry (Sharif HPC Cluster)

• Physical Compute: NVIDIA GeForce RTX 3090 24GB (Ampere Architecture, Tensor Cores active)
• Training Duration: 49.0s
• Total Optimization Steps: 118 steps (Stage 1 Subspace Hidden Alignment $\rightarrow$ Stage 2 Logit KL + Task CE)
• Effective Throughput: 19739.7 tokens/sec
• Peak GPU VRAM: 13512 MB (Zero teacher VRAM overhead via offline representation caching)

---

🚀 Quick Start & One-Liner Inference

1. Standard Hugging Face Transformers (AutoModelForSequenceClassification)

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

# Instant download and execution directly from Hugging Face Hub (trust_remote_code=True)
model = AutoModelForSequenceClassification.from_pretrained("tahamajs/bert-factorised-matrix-state-sst2", trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained("tahamajs/bert-factorised-matrix-state-sst2")

inputs = tokenizer("NeuralOps accelerates transformer architectures with subquadratic efficiency.", return_tensors="pt")
with torch.no_grad():
    outputs = model(**inputs)
    prediction = outputs.logits.argmax(dim=-1).item()

print("Predicted Sentiment (0 = Negative, 1 = Positive):", prediction)

2. High-Level Modular Loading via NeuralOps Zoo API

from neuralops.hub.bert_zoo import load_pretrained_bert

# Reconstruct and load strictly the attention weights onto base BERT
model, tokenizer = load_pretrained_bert(
    attention_type="factorised_matrix_state",
    task="sequence-classification",
    hf_org="tahamajs",
)

---

📦 Modular Attention-Only Weights

Following the NeuralOps parameter-efficient serialization specification, this repository stores only the specialized attention and task adapter parameters (attention_weights.safetensors in FP16, ~55 MB) rather than re-uploading the entire 450 MB frozen base BERT weights. The remote code shim (modeling_neuralops_bert.py) automatically initializes the frozen base weights from bert-base-uncased and applies the distilled attention parameters on the fly.

---

📜 Citation

If you use this model or the NeuralOps attention zoo in your research, please cite:

@software{neuralops2026bertzoo,
  title   = {NeuralOps: Pretrained BERT Subquadratic Attention Zoo on GLUE Benchmark},
  author  = {Majlesi, Taha},
  year    = {2026},
  url     = {https://github.com/tahamajs/NeuralOps}
}