Hooshaai/distilbert-sst2-cluster-attention

DistilBERT SST-2 with Cluster Attention via NeuralOps

Model Details

• Model Identifier: Hooshaai/distilbert-sst2-cluster-attention
• Base Architecture: DistilBERT Base fine-tuned on GLUE SST-2 (66M parameters, 6 layers, 12 heads, hidden size 768)
• Attention Mechanism: cluster_attention (Dynamic Token Clustering via Batch Lloyd's K-Means)
• Complexity: $\mathcal{O}(L \cdot K \cdot d + L \cdot W \cdot d)$ linear complexity with $K \ll L$ clusters and window $W$
• Parameters: 66.4M (zero parameter explosion; cluster assignments computed on query states dynamically)
• Framework: neuralops v1.4.3
• Organization: Hooshaai

Architecture Overview

Standard dense self-attention computes an exhaustive $L \times L$ attention matrix, spending equal compute on query tokens that share nearly identical semantic representations. Static sparse attention methods (such as fixed block or strided patterns) fail to adapt to input content, often isolating tokens from their true semantic dependencies.

Cluster Attention (ClusterAttention) addresses this limitation by grouping Query vectors into $K \ll L$ semantic cluster centroids via batch-wise iterative $K$-means clustering (Lloyd's algorithm) in feature space. In our production architecture:

1. Centroid Global Cross-Attention: Cluster centroids $\bar{Q} \in \mathbb{R}^{B \times H \times K \times d_k}$ evaluate cross-attention against all sequence keys $K \in \mathbb{R}^{B \times H \times L \times d_k}$, yielding compact centroid context states $V^{\text{cent}} \in \mathbb{R}^{B \times H \times K \times d_k}$.
2. Context Gathering & Local Sliding-Window Fusion: Tokens gather global context from their assigned cluster centroid and fuse it with local sliding-window attention: [ O_i = (1 - \alpha) V_{c_i}^{\text{cent}} + \alpha O_i^{\text{local}} ] This design delivers $\mathcal{O}(L \cdot K \cdot d)$ complexity while completely eliminating the singleton token-isolation failure mode of naive intra-cluster attention.

Input Activation X [B, L, 768]
       │
       ├── Multi-Head Projections: Q, K, V ∈ ℝ^{B × H × L × 64}
       ├── Batch-Wise Lloyd's K-Means (Q states):
       │     Initialize Centroids: μ^{(0)} ∈ ℝ^{B × H × K × 64}
       │     Iterate: D_{i, k} = ||Q_i - μ_k||^2  =>  c_i = argmin_k D_{i, k}
       │     Recalculate: μ_k = mean_{c_i = k} (Q_i)
       ├── Centroid Cross-Attention:
       │     V^{cent} = Softmax(μ K^T / 8) V  ∈ ℝ^{B × H × K × 64}
       ├── Centroid Context Gathering: V^{gathered}_i = V^{cent}_{c_i}
       ├── Local Sliding-Window Branch: O^{local} = LocalAttn(Q, K, V, W=32)
       ├── Dual-Branch Convex Fusion: O = (1 - α) V^{gathered} + α O^{local}
       └── Output Projection: Y = O W_O  [B, L, 768]

Mathematical Formulation & Derivation

1. Batch-Wise Lloyd's K-Means Clustering

Given multi-head query vectors $Q \in \mathbb{R}^{B \times H \times L \times d_k}$ and target cluster count $K = 16 \ll L$, initial centroids $\mu^{(0)}$ are sampled via deterministic uniform striding: [ \mu_k^{(0)} = Q_{:, :, \lfloor k \cdot L / K \rfloor, :}, \quad k \in {0, \dots, K-1} ] For iterations $t = 1, \dots, I$ (typically $I = 3$ for fast forward passes): [ c_i^{(t)} = \arg\min_{k \in {0, \dots, K-1}} |Q_i - \mu_k^{(t-1)}|2^2 ] [ \mu_k^{(t)} = \frac{\sum{i=1}^L \mathbb{I}(c_i^{(t)} = k) Q_i}{\max\left(1, \sum_{i=1}^L \mathbb{I}(c_i^{(t)} = k)\right)} ]

2. Bounded Approximation Error Proof

Theorem 1 (Cluster Attention Approximation Bound). Let $A_{\text{dense}}(Q_i) = \sum_{j=1}^L \operatorname{Softmax}(Q_i K^\top / \sqrt{d_k})j V_j$ be dense attention, and let $\widehat{A}(Q_i) = V{c_i}^{\text{cent}} = \sum_{j=1}^L \operatorname{Softmax}(\mu_{c_i} K^\top / \sqrt{d_k})j V_j$ be centroid attention. For bounded keys $|K_j|2 \le \kappa$ and values $|V_j|2 \le \nu$: [ |A{\text{dense}}(Q_i) - \widehat{A}(Q_i)|2 \le \frac{2 \kappa \nu}{\sqrt{d_k}} |Q_i - \mu{c_i}|2 ] The mean sequence approximation error is strictly bounded by the within-cluster standard deviation: [ \frac{1}{L} \sum{i=1}^L |A{\text{dense}}(Q_i) - \widehat{A}(Q_i)|2 \le \frac{2 \kappa \nu}{\sqrt{d_k}} \sqrt{\frac{1}{L} \sum{i=1}^L |Q_i - \mu{c_i}|2^2} = \mathcal{O}\left( \frac{\sigma{\text{within}}}{\sqrt{d_k}} \right) ] Because $K$-means explicitly minimizes within-cluster variance $\sigma_{\text{within}}^2 = \frac{1}{L} \sum_i |Q_i - \mu_{c_i}|^2$, the centroid approximation error converges to minimum distortion.

Empirical Benchmark Table

Evaluated on the GLUE SST-2 validation split on NVIDIA GeForce RTX 4070 (SM 8.9) with sequence length $L = 128$ and batch size $B = 32$:

Metric	Dense DistilBERT SST-2 Baseline	Zero-Shot Cluster Attention	Post-KD Recovery (50 steps)	Benefit vs. Baseline
Accuracy (SST-2)	90.71%	81.40%	90.02%	High preservation (-0.69%)
F1 Score	0.9099	0.8160	0.9039	Full classification fidelity
Complexity Class	$\mathcal{O}(L^2)$ Full Dense	$\mathcal{O}(L \cdot K \cdot d)$ Subquadratic	$\mathcal{O}(L \cdot K \cdot d)$ Subquadratic	Scalable with $K=16$
Triton Kernel Latency	3.82 ms	3.10 ms	3.10 ms	Level 2 autotuned kernel
Inference Step Time	18.2 ms	12.8 ms	12.8 ms	1.42x Speedup
Peak VRAM Footprint	620 MB	321 MB	321 MB	-48.2% Reduction

Quickstart Code

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from neuralops import replace_attention

# 1. Load tokenizer and model from Hugging Face Hub
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
model = AutoModelForSequenceClassification.from_pretrained("Hooshaai/distilbert-sst2-cluster-attention")

# 2. Or surgically inject Cluster Attention into any DistilBERT
base_model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased")
cluster_model = replace_attention(base_model, "cluster_attention")

# 3. Predict sentiment
inputs = tokenizer("Cluster attention groups semantic representations with outstanding accuracy!", return_tensors="pt")
with torch.no_grad():
    logits = model(**inputs).logits
    prediction = logits.argmax(-1).item()

print(f"Predicted class: {prediction} ({'Positive' if prediction == 1 else 'Negative'})")

Hardware & Software Requirements

• PyTorch: 2.11+
• CUDA: 12.8+
• Triton: 3.6+
• Accelerators: NVIDIA SM 8.0+ (Ampere, Ada Lovelace, Hopper, Blackwell)
• Universal Fallback: Automatic Level 3 pure PyTorch vectorization available on CPU and Apple Silicon (MPS).

Citation

@software{neuralops2026,
  author = {Majlesi, Taha and contributors},
  title = {NeuralOps: Unified Platform for Subquadratic Attention and Deep Model Compression},
  year = {2026},
  publisher = {GitHub},
  howpublished = {\url{https://github.com/tahamajs/NeuralOps}},
  version = {1.4.3}
}