DistilBERT SST-2 with RPC Attention via NeuralOps

Model Details

• Model Identifier: Hooshaai/distilbert-sst2-rpc
• Base Architecture: DistilBERT Base fine-tuned on GLUE SST-2 (66M parameters, 6 layers, 12 heads, hidden size 768)
• Attention Mechanism: rpc_attention (Randomized Phase Cluster / Random Projection Compressed Attention)
• Complexity: Compressed scoring FLOPs with $d_p \ll d_k$ subspace projection while retaining full value context
• Parameters: 66.4M (fixed non-trainable orthogonal sketch matrix $P$, zero parameter increase)
• Framework: neuralops v1.4.3
• Organization: Hooshaai

Architecture Overview

Standard multi-head attention computes pairwise inner products between queries and keys in the full head dimension $d_k = 64$. For long sequences or latency-sensitive inference, evaluating full $d_k$-dimensional inner products consumes significant memory bandwidth and arithmetic units.

Random Projection Compressed (RPC) Attention compresses query-key inner product evaluations by projecting head states into a lower-dimensional subspace $d_p = 32$ ($d_p \ll d_k$).

RPC Attention leverages the Johnson-Lindenstrauss (JL) Lemma, which guarantees that pairwise Euclidean distances and cosine inner products between vectors in high-dimensional spaces are strictly preserved under random orthogonal projections. By initializing fixed, non-trainable head-wise random projection matrices $P \in \mathbb{R}^{H \times d_k \times d_p}$, RPC Attention computes attention score logits in compressed subspace $\mathbb{R}^{d_p}$, cutting scoring FLOPs by 50% and dramatically accelerating kernel execution. Full-dimensional values $V \in \mathbb{R}^{d_k}$ are preserved to ensure representation expressivity.

Input Representation X [B, L, 768]
       │
       ├── Multi-Head Projections: Q, K, V ∈ ℝ^{B × H × L × 64}
       │
       ├── Fixed Johnson-Lindenstrauss Subspace Projection P ∈ ℝ^{H × 64 × 32}:
       │     Q_sketch = Q · P  ∈ ℝ^{B × H × L × 32}
       │     K_sketch = K · P  ∈ ℝ^{B × H × L × 32}
       │
       ├── Sketched Attention Logits:
       │     S = (Q_sketch K_sketch^T) / sqrt(32)  ∈ ℝ^{B × H × L × L}
       │     A = Softmax(S)
       │
       ├── Uncompressed Value Context Aggregation:
       │     O = A V  ∈ ℝ^{B × H × L × 64}
       │
       └── Output Projection: Y = Concat(O_1, ..., O_H) W_O  [B, L, 768]

Mathematical Formulation & Derivation

1. Johnson-Lindenstrauss Random Subspace Sketch

Let $P_h \in \mathbb{R}^{d_k \times d_p}$ be an orthogonal Gaussian random projection matrix with entries independently sampled: [ P_{h, i, j} \sim \mathcal{N}\left( 0, ; \frac{1}{d_p} \right) ] Projecting query vectors $q_i \in \mathbb{R}^{d_k}$ and key vectors $k_j \in \mathbb{R}^{d_k}$: [ \tilde{q}_i = q_i P_h \in \mathbb{R}^{d_p}, \quad \tilde{k}_j = k_j P_h \in \mathbb{R}^{d_p} ] By the Johnson-Lindenstrauss Lemma: [ (1 - \epsilon) |q_i - k_j|_2^2 \le |\tilde{q}_i - \tilde{k}_j|_2^2 \le (1 + \epsilon) |q_i - k_j|_2^2 ] with probability $\ge 1 - 2 e^{-(\epsilon^2 - \epsilon^3) d_p / 4}$.

2. Inner Product Preservation

Expanding the norm relation: [ |\tilde{q}_i - \tilde{k}_j|_2^2 = |\tilde{q}_i|_2^2 + |\tilde{k}_j|_2^2 - 2 \langle \tilde{q}_i, \tilde{k}_j \rangle ] Because random Gaussian projections preserve vector lengths ($\mathbb{E}[|\tilde{q}|_2^2] = |q|_2^2$), the inner product is an unbiased estimator: [ \mathbb{E}\left[ \langle \tilde{q}_i, \tilde{k}_j \rangle \right] = \langle q_i, k_j \rangle ] with variance $\operatorname{Var}(\langle \tilde{q}_i, \tilde{k}_j \rangle) = \frac{|q_i|^2 |k_j|^2 + \langle q_i, k_j \rangle^2}{d_p} = \mathcal{O}\left( \frac{1}{d_p} \right)$.

3. Sketched Attention and Full Value Aggregation

Attention scores are computed in the compressed subspace: [ S_{h, i, j} = \frac{\tilde{q}{h, i} \tilde{k}{h, j}^\top}{\sqrt{d_p}} + M_{\text{mask}} ] [ A_{h, i, j} = \operatorname{Softmax}j\left( S{h, i, j} \right) ] Values $V_h \in \mathbb{R}^{L \times d_k}$ remain in their uncompressed feature space: [ O_h = A_h V_h \in \mathbb{R}^{L \times d_k} ]

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$:

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

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

# 3. Predict sentiment
inputs = tokenizer("Random subspace sketching delivers rapid attention scoring with theoretical bounds!", 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}
}