Hooshaai/distilbert-sst2-adaptive-span

DistilBERT SST-2 with Adaptive Span Attention via NeuralOps

Model Details

• Model Identifier: Hooshaai/distilbert-sst2-adaptive-span
• Base Architecture: DistilBERT Base fine-tuned on GLUE SST-2 (66M parameters, 6 layers, 12 heads, hidden size 768)
• Attention Mechanism: adaptive_span_attention (Dynamic Head-Specific Context Window Gating)
• Complexity: Dynamic subquadratic time and memory complexity bounded by $\mathcal{O}(L \cdot \bar{S})$ where $\bar{S} \ll L$ is the mean active span
• Parameters: 66.4M (+12 learnable span logits $\theta_h$, parameter overhead < 0.001%)
• Framework: neuralops v1.4.3
• Organization: Hooshaai

Architecture Overview

In standard transformer multi-head self-attention, all attention heads compute affinities across the entire context window $L$, forcing heads that specialize in local syntactic agreements to compute quadratic interactions with distant, irrelevant tokens.

Adaptive Span Attention endows each attention head $h \in {1, \dots, H}$ with an independent, continuously differentiable receptive field size $S_h \in [0, S_{\max}]$. Heads that resolve local syntax learn compact spans ($S_h \approx 16-32$), while long-range semantic heads maintain wider spans ($S_h \approx 128$). Tokens positioned beyond a head's effective span boundary are strictly masked out. By pairing continuous soft-ramp masking during distillation with discrete boundary clamping during deployment, Adaptive Span Attention eliminates redundant computation and reduces VRAM footprint by 48.2% while retaining 90.94% SST-2 accuracy.

Input Activation X [B, L, 768]
       │
       ├── Multi-Head Projections: Q, K, V ∈ ℝ^{B × H × L × 64}
       ├── Head Span Evaluation: S_h = σ(θ_h) · S_max ∈ [0, S_max]
       ├── Relative Distance Matrix: Δ_{i, j} = |i - j|
       ├── Differentiable Ramp Masking: M_{h, i, j} = clamp((R + S_h - Δ_{i, j}) / R, 0, 1)
       ├── Masked Dot-Product Attention:
       │     A_{h, i, j} = Softmax_j( (Q_h K_h^T / 8) + log(M_{h, i, j} + ε) )
       └── Context Synthesis: Y = Concat(A_1 V_1, ..., A_H V_H) W_O [B, L, 768]

Mathematical Formulation & Derivation

1. Learnable Continuous Span Parameterization

For each attention head $h \in {1, \dots, H}$, let $\theta_h \in \mathbb{R}$ denote a learnable logit parameter. The effective continuous attention span $S_h$ is defined via Sigmoid gating: [ S_h = S_{\max} \cdot \sigma(\theta_h) = \frac{S_{\max}}{1 + \exp(-\theta_h)} \in [0, S_{\max}] ] where $S_{\max}$ is the maximum sequence length (e.g. 128 or 512).

2. Differentiable Soft-Ramp Masking Function

To allow backpropagation into discrete window boundaries during recovery distillation, a soft-ramp function $g_{S_h, R}(d)$ with ramp buffer width $R > 0$ modulates the relative token distance $d = |i - j|$: [ g_{S_h, R}(d) = \operatorname{clip}\left( \frac{1}{R} (R + S_h - d), ; 0, ; 1 \right) = \begin{cases} 1 & \text{if } d \le S_h \ \frac{R + S_h - d}{R} & \text{if } S_h < d \le S_h + R \ 0 & \text{if } d > S_h + R \end{cases} ] The attention scoring matrix $\widehat{S}{h, i, j}$ incorporates the ramp mask: [ \widehat{S}{h, i, j} = \frac{Q_{h, i} K_{h, j}^\top}{\sqrt{d_k}} + \log(g_{S_h, R}(|i - j|) + \epsilon) ] and the normalized attention probability evaluates as: [ A_{h, i, j} = \frac{\exp(\widehat{S}{h, i, j})}{\sum{l=1}^L \exp(\widehat{S}_{h, i, l})} ] In inference mode, tokens with $d > S_h$ are strictly skipped in Level 1 CUDA C++ and Level 2 Triton kernels, yielding subquadratic scaling $\mathcal{O}(L \cdot \bar{S})$.

Empirical Benchmark Table

Evaluated on the GLUE SST-2 validation benchmark 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 Adaptive Span	Post-Distillation (v1.4.3)	Change vs. Baseline
Accuracy (SST-2)	90.71%	83.20%	90.94%	+0.23% (Net Gain)
F1 Score	0.9099	0.8350	0.9033	-0.0066
Complexity Class	$\mathcal{O}(L^2)$ Full Dense	$\mathcal{O}(L \cdot \bar{S})$ Subquadratic	$\mathcal{O}(L \cdot \bar{S})$ Subquadratic	Dynamic Windowing
Triton Kernel Latency	3.82 ms	1.24 ms	1.24 ms	3.08x Speedup
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
Average Head Span ($\bar{S}$)	128 (100%)	128 (Initial)	58.4 (45.6%)	54.4% Sparsified

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-adaptive-span")

# 2. Or dynamically inject Adaptive Span Attention into any DistilBERT
base_model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased")
adaptive_model = replace_attention(base_model, "adaptive_span_attention")

# 3. Predict sentiment
inputs = tokenizer("The adaptive span mechanisms dynamically prune context with zero accuracy loss!", 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}
}