Hooshaai/distilbert-mnli-full-adasplash

DistilBERT MNLI with AdaSplash Attention via NeuralOps

Model Details

• Model Identifier: Hooshaai/distilbert-mnli-full-adasplash
• Base Architecture: DistilBERT Base fine-tuned on Multi-Genre Natural Language Inference (MNLI, 3-class classification, 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 zero weights ($A_{ij} \equiv 0$) across multi-sentence premise-hypothesis pairs
• Parameters: 67.5M (+0.88% parameter count via AdaSplashLinear residual LoRA adapters)
• Framework: neuralops v1.4.3
• Organization: Hooshaai

Architecture Overview

Natural Language Inference (NLI) pairs complex premise and hypothesis sentences, requiring the model to determine whether the hypothesis is an entailment, contradiction, or neutral with respect to the premise. Because premise sentences often contain subordinate clauses and verbose details, standard dense Softmax attention tends to spread attention mass diffusely across irrelevant words, degrading discriminative inference.

AdaSplash Attention resolves diffuse leakage by projecting attention affinities onto the probability simplex using Tsallis $\alpha = 1.5$ non-extensive entmax:

1. Precise Cross-Sentence Sparsification: Eliminates non-informative token interactions by assigning exactly zero probability mass ($A_{ij} \equiv 0$) to roughly 42% of token pairs.
2. Stable Numerical Lagrange Multiplier Finding: Solves the normalization constraint $\sum [ (s_i - \tau)/2 ]_+^2 = 1$ via bounded bisection search in at most 25 iterations.
3. AdaSplashLinear Residual Adapters: Pairs frozen pre-trained projection matrices with trainable low-rank adapters ($W = W_{\text{base}} + B \cdot A$), maintaining 81.15% MNLI accuracy (preserving 99.2% of baseline accuracy).

Premise & Hypothesis Pair X [B, L, 768]
       │
       ├── AdaSplashLinear Projections: Q, K, V via (W_base + B · A)
       ├── Cross-Sentence Affinity Matrix: 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 Entmax Output: A_{i, j} = max(0, (S_{i, j} - τ)/2)^2
       │     => 42% exact zeros suppress irrelevant background cross-attention
       ├── Context Synthesis: 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 Simplex Formulation

For probability vector $p \in \Delta^{L-1}$ and $\alpha = 1.5$: [ H_{1.5}(p) = \frac{4}{3} \sum_{i=1}^L \left( p_i - p_i^{1.5} \right) ] The 1.5-entmax projection evaluates: [ \operatorname{entmax}{1.5}(s) = \arg\max{p \in \Delta^{L-1}} \left{ p^\top s + \frac{4}{3} \sum_{i=1}^L (p_i - p_i^{1.5}) \right} ]

2. Analytical Thresholding

The closed-form solution satisfies: [ p_i^*(s) = \left[ \frac{s_i - \tau}{2} \right]+^2 ] where Lagrange multiplier $\tau$ is chosen such that $\sum{i=1}^L p_i^*(s) = 1$. When token $j$ has logit $s_j \le \tau$, its attention weight is identically zero: [ p_j^*(s) = 0, \quad \forall j \text{ s.t. } s_j \le \tau ] In cross-sentence reasoning, tokens in the hypothesis exclusively attend to informative tokens in the premise while completely pruning out filler tokens and distant syntax.

Empirical Benchmark Table

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

Metric	Dense DistilBERT MNLI Baseline	Zero-Shot AdaSplash	Post-Distillation (v1.4.3)	Benefit vs. Baseline
Accuracy (MNLI Matched)	81.80%	76.40%	81.15%	Minimal delta (-0.65%)
F1 Macro	0.8175	0.7630	0.8110	High 3-class precision
Attention Sparsity (% Zeros)	0.0% (Softmax)	42.0%	42.0%	42% Hard Zero 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-mnli-full-adasplash")

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

# 3. Predict inference relation (0: Contradiction, 1: Entailment, 2: Neutral)
premise = "The experimental algorithm demonstrates subquadratic scaling on all benchmarks."
hypothesis = "The algorithm runs faster than quadratic self-attention."

inputs = tokenizer(premise, hypothesis, return_tensors="pt")
with torch.no_grad():
    logits = model(**inputs).logits
    prediction = logits.argmax(-1).item()

labels = ["Contradiction", "Entailment", "Neutral"]
print(f"Predicted relation: {labels[prediction]} (Class {prediction})")

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}
}