DistilBERT SST-2 with AdaSplash Attention via NeuralOps

Model Details

  • Model Identifier: Hooshaai/distilbert-sst2-adasplash
  • Base Architecture: DistilBERT Base fine-tuned on GLUE SST-2 (66M parameters, 6 layers, 12 heads, hidden size 768)
  • Attention Mechanism: adasplash_attention (Adaptive Splash Attention with Tsallis $\alpha=1.5$ Entmax Sparse Simplex Projection)
  • Attention Sparsity: 42% exact mathematical zero weights ($A_{ij} \equiv 0$)
  • Parameters: 67.5M (+0.88% parameter count via AdaSplashLinear residual LoRA adapters)
  • Framework: neuralops v1.4.3
  • Organization: Hooshaai

Architecture Overview

Standard transformer attention relies on Softmax activation, which maps unbounded dot-product logits into strictly positive probabilities ($\exp(x) > 0, \forall x \in \mathbb{R}$). As a consequence, Softmax forces every single token in the sequence to assign non-zero probability mass to every other token. This causes diffuse attention leakage, where background tokens, irrelevant punctuation, and distant padding contaminate the contextual embedding of salient tokens.

AdaSplash Attention introduces an adaptive, mathematically rigorous sparse attention mechanism combining:

  1. Exact Mathematical Sparsity (1.5-Entmax): By replacing Shannon entropy with Tsallis non-extensive entropy ($\alpha = 1.5$), AdaSplash projects logits onto the probability simplex such that irrelevant tokens receive identically zero attention weight ($A_{ij} \equiv 0$), resulting in 42% exact zeros in DistilBERT.
  2. Numerically Stabilized Bisection Root Finding: Computes the exact Lagrange normalizer $\tau$ in at most 25 iterations with error $\le 2^{-25} \approx 3 \times 10^{-8}$, eliminating NaN gradients in FP16/BF16.
  3. AdaSplashLinear with Low-Rank Residual Adaptation: Projections ($W_Q, W_K, W_V, W_O$) encapsulate frozen pre-trained weights augmented with trainable LoRA adapters ($W = W_{\text{base}} + B \cdot A$), preserving pre-trained knowledge while recovering downstream accuracy to 90.94% (+0.23% over baseline).
  4. Bidirectional Full-Sequence Splash: While autoregressive causal decoders apply local splash windowing ($W=32$), bidirectional encoders retain full global reach with exact entmax sparsity to prevent context truncation.
Input Representation X [B, L, 768]
       β”‚
       β”œβ”€β”€ AdaSplashLinear Projections: Q, K, V via (W_base + B Β· A)
       β”œβ”€β”€ Unnormalized Affinities: S_{i, j} = (Q_i K_j^T) / 8
       β”œβ”€β”€ Exact Tsallis 1.5-Entmax Simplex Projection:
       β”‚     Bisection Search: find unique Ο„ such that βˆ‘ max(0, (S_{i, j} - Ο„)/2)^2 = 1
       β”‚     Sparse Routing Matrix: A_{i, j} = max(0, (S_{i, j} - Ο„)/2)^2
       β”‚     => 42% of entries are EXACT ZERO (A_{i, j} ≑ 0)
       β”œβ”€β”€ Sparse Context Aggregation: O_i = βˆ‘_{j: A_{i, j} > 0} A_{i, j} V_j
       └── Output Projection: Y = AdaSplashLinear_O(O)  [B, L, 768]

Mathematical Formulation & Derivation

1. Tsallis $\alpha$-Entropy and Simplex Projection

Let $\Delta^{L-1} = { p \in \mathbb{R}+^L : \sum{i=1}^L p_i = 1 }$ denote the probability simplex in $\mathbb{R}^L$. For $\alpha = 1.5$, the Tsallis non-extensive entropy $H_{1.5}(p)$ is: [ H_{1.5}(p) = \frac{1}{1.5 \cdot 0.5} \sum_{i=1}^L \left( p_i - p_i^{1.5} \right) = \frac{4}{3} \sum_{i=1}^L \left( p_i - p_i^{1.5} \right) ] The 1.5-entmax mapping is the regularized argmax: [ \operatorname{entmax}{1.5}(s) = \arg\max{p \in \Delta^{L-1}} \left{ p^\top s + H_{1.5}(p) \right} ]

2. Closed-Form Simplex Projection

Theorem 1 (1.5-Entmax Closed Form). For an arbitrary logit vector $s \in \mathbb{R}^L$, the unique optimal probability vector $p^* = \operatorname{entmax}{1.5}(s)$ is: [ p_i^*(s) = \left[ \frac{s_i - \tau}{2} \right]+^2 = \max\left(0, ; \frac{s_i - \tau}{2}\right)^2 ] where $[x]+ = \max(0, x)$ and $\tau \in \mathbb{R}$ is the unique Lagrange multiplier satisfying: [ g(\tau) \triangleq \sum{i=1}^L \left[ \frac{s_i - \tau}{2} \right]_+^2 = 1 ] Because $g(\tau)$ is strictly monotonically decreasing on $(-\infty, \max(s))$, the root $\tau^*$ is found via bisection search within bounds $[\max(s) - 2, \max(s) - 2/\sqrt{L}]$ in at most 25 iterations.

3. Exact Sparsity Support

Tokens with logits $s_j \le \tau$ satisfy: [ p_j^*(s) \equiv 0 ] This enforces a hard structural threshold below which tokens exert zero influence, eliminating background noise contamination.

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 AdaSplash Post-Distillation (v1.4.3) Benefit vs. Baseline
Accuracy (SST-2) 90.71% 77.18% 90.94% +0.23% (Net Gain)
F1 Score 0.9099 0.7750 0.9125 +0.0026
Attention Sparsity (% Zeros) 0.0% (Softmax) 42.0% 42.0% Strict Mathematical Sparsity
Active Parameters 66.4M 67.5M 67.5M +0.88% LoRA adapters
Triton Kernel Latency 3.82 ms 1.54 ms 1.54 ms 2.48x Kernel 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

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-adasplash")

# 2. Or surgically inject AdaSplash Attention into any DistilBERT in 1 line
base_model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased")
adasplash_model = replace_attention(base_model, "adasplash_attention")

# 3. Predict sentiment
inputs = tokenizer("AdaSplash exact sparsity eliminates diffuse noise and sharpens predictions!", 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}
}
Downloads last month

-

Downloads are not tracked for this model. How to track
Inference Providers NEW
This model isn't deployed by any Inference Provider. πŸ™‹ Ask for provider support

Model tree for Hooshaai/distilbert-sst2-adasplash

Finetuned
(12754)
this model

Datasets used to train Hooshaai/distilbert-sst2-adasplash