Sparsefinder: Efficient Sparse Entmax

• Sparsefinder is a method that predicts the exact sparsity pattern of entmax-based attention, reducing dense matrix computation by focusing on nonzero entries.
• It uses low-dimensional projections with contrastive hinge loss and grouping strategies (distance, quantization, clustering) to efficiently approximate attention edges.
• Empirical evaluations in machine translation and language modeling show that Sparsefinder achieves near full recall with significant reductions in FLOPs and memory usage.

Sparsefinder is a method for efficiently predicting the exact sparsity pattern of entmax-based attention in transformer models, allowing one to compute only those attention edges that will receive nonzero probability, with the goal of significant reductions in both computation and memory. Originally introduced by Treviso et al. (2021), Sparsefinder addresses the challenge that, while entmax produces sparse attention distributions, it still requires forming and normalizing a dense attention matrix with $O(nm)$ complexity, where $n$ and $m$ are the sequence lengths. Sparsefinder predicts the sparse attention graph prior to the actual attention computation, enabling focused computation on attended positions and achieving improved efficiency without sacrificing the accuracy or recall of the attention mechanism (Treviso et al., 2021).

1. Mathematical Foundations of Entmax Attention

Sparsefinder builds directly on entmax$_\alpha$ attention mechanisms. Whereas standard Transformer attention computes $Z = QK^\top \in \mathbb{R}^{n\times m}$, with queries $Q \in \mathbb{R}^{n\times d}$ and keys $K \in \mathbb{R}^{m\times d}$, then applies softmax normalization giving dense distributions, entmax$_\alpha$ replaces softmax with a Tsallis-type entropy regularizer. The entmax$_\alpha$ function is defined by

$$\text{entmax}_\alpha(z) = \arg\max_{p\in\Delta^m} \left( p^\top z - H_\alpha(p) \right)$$

where $n$0 for $n$1 (recovering softmax), and for $n$2,

$n$3

For $n$4, the attention distribution is sparse: coordinates with $n$5 are exactly zero, and nonzero values have the closed-form

$n$6

Popular choices include $n$7 (sparsemax) and $n$8 (typical in this work for a trade-off between sparsity and ease of computation).

The "attention graph" for each head is

$n$9

which is typically much sparser than the dense $m$0 matrix. Due to the "sparse-consistency" property of entmax, if a mask $m$1 is predicted and the attention computation is restricted to this mask, the result is exactly equal on all truly nonzero entries.

2. Sparsefinder Model Architecture and Variants

Sparsefinder operates in two key stages:

(A) Learning Low-Dimensional Projections:

For every head $m$2, learn linear projections $m$3 (with $m$4) and compute projected queries and keys $m$5, $m$6. Training is by a contrastive hinge loss:

$m$7

where $m$8 is a margin. The result is a space in which true attended pairs are close.

(B) Edge Prediction via Grouping Strategies:

• Distance-Based Sparsefinder:

Predict $m$9 as an edge if $_\alpha$0 for threshold $_\alpha$1. All $_\alpha$2 distances are computed, trading off sparsity and recall by tuning $_\alpha$3.

• Quantization-Based Sparsefinder:

Each of $_\alpha$4 dimensions is partitioned into $_\alpha$5 quantile bins. Tokens sharing a bin in at least one dimension are assumed adjacent. No additional learning is needed beyond projection. Edges scale as $_\alpha$6.

• Clustering-Based Sparsefinder:

Offline $_\alpha$7-means clustering is run on projected vectors, yielding $_\alpha$8 centroids. Each $_\alpha$9, $Z = QK^\top \in \mathbb{R}^{n\times m}$0 is assigned to its $Z = QK^\top \in \mathbb{R}^{n\times m}$1 nearest centroids. Edge $Z = QK^\top \in \mathbb{R}^{n\times m}$2 is predicted present if assigned clusters intersect. The parameter $Z = QK^\top \in \mathbb{R}^{n\times m}$3 controls sparsity/recall.

Sparsefinder can be combined with fixed local patterns (e.g., sliding windows of width $Z = QK^\top \in \mathbb{R}^{n\times m}$4) and global tokens as used in architectures such as Longformer or BigBird.

3. Computational and Memory Complexity

The computational improvements arise from predicting $Z = QK^\top \in \mathbb{R}^{n\times m}$5 so that full $Z = QK^\top \in \mathbb{R}^{n\times m}$6 need not be formed.

Approach	FLOPs (score computation)	Memory (mask)
Dense entmax$Z = QK^\top \in \mathbb{R}^{n\times m}$7	$Z = QK^\top \in \mathbb{R}^{n\times m}$8	$Z = QK^\top \in \mathbb{R}^{n\times m}$9
Sparsefinder-projection	$Q \in \mathbb{R}^{n\times d}$0	$Q \in \mathbb{R}^{n\times d}$1
Sparsefinder-mask (dist.)	$Q \in \mathbb{R}^{n\times d}$2	$Q \in \mathbb{R}^{n\times d}$3
Sparsefinder-dot products	$Q \in \mathbb{R}^{n\times d}$4 ($Q \in \mathbb{R}^{n\times d}$5)	--

For quantization, bucket assignment cost is $Q \in \mathbb{R}^{n\times d}$6 and candidate edges $Q \in \mathbb{R}^{n\times d}$7. Clustering-based mask generation is $Q \in \mathbb{R}^{n\times d}$8 and candidate edges scale as $Q \in \mathbb{R}^{n\times d}$9 for balanced clusters with $K \in \mathbb{R}^{m\times d}$0 assignments per token. If $K \in \mathbb{R}^{m\times d}$1 and $K \in \mathbb{R}^{m\times d}$2, substantial savings are obtained relative to the standard $K \in \mathbb{R}^{m\times d}$3 complexity.

4. Empirical Evaluation in Machine Translation and Language Modeling

Sparsefinder was evaluated in both machine translation (MT) and masked language modeling (MLM) settings:

MT Experiments:

A "transformer-large" (6 layers, 16 heads, $K \in \mathbb{R}^{m\times d}$4) model was pretrained on Paracrawl and fine-tuned with $K \in \mathbb{R}^{m\times d}$5-entmax ($K \in \mathbb{R}^{m\times d}$6) on IWSLT'17 EN$K \in \mathbb{R}^{m\times d}$7DE and EN$K \in \mathbb{R}^{m\times d}$8FR, with BPE vocab size 32k. Projections from $K \in \mathbb{R}^{m\times d}$9 to $_\alpha$0 dimensions were learned per head, using one epoch of margin-hinge loss on 10k long sentences ($_\alpha$18--9M positive pairs/layer). Hyperparameters were swept over distance thresholds $_\alpha$2, bucket counts $_\alpha$3, and window sizes $_\alpha$4.

• Pareto frontiers of sparsity vs. recall showed that distance-based and clustering-based Sparsefinder variants strictly dominate baselines (Longformer, BigBird, Reformer, Routing Transformer, local window) for any given sparsity
• For both EN$_\alpha$5DE and EN$_\alpha$6FR, distance and clustering recover nearly full BLEU at the sparsity of the dense entmax Transformer, e.g., BLEU 34.47 (EN$_\alpha$7DE), 42.65 (EN$_\alpha$8FR)

MLM Experiments:

RoBERTa-base ($_\alpha$9, $_\alpha$0) was fine-tuned on WikiText-103 for 3k steps with $_\alpha$1 entmax ($_\alpha$2, overall sparsity $_\alpha$3). Projection from $_\alpha$4 to $_\alpha$5 was trained, and grouping parameters swept.

• Distance and clustering Sparsefinder achieved strictly better recall or perplexity for all sparsity levels vs. Longformer, BigBird, Reformer, Routing Transformer, and local baselines

All experiments ran on NVIDIA Titan Xp/GTX 1080Ti/RTX 2080 Ti GPUs and AMD/Intel CPUs; projection learning used 128 positive/negative pairs, Adam LR=$_\alpha$6, margin $_\alpha$7, for one epoch.

5. Design Trade-Offs and Practical Considerations

The primary controlling parameters are the distance threshold $_\alpha$8 (distance-based) and cluster-assignment $_\alpha$9 (clustering-based), directly trading sparsity for recall: lower values increase sparsity while reducing recall; higher values increase recall at the expense of more computation. Quantization-based grouping, while algorithmically simplest, underperforms: it cannot adapt to irregular neighborhoods in the projection space.

Clustering achieves performance near that of distance-based grouping but at lower computational cost, since pairwise distances are not required. For balanced clusters, $$\text{entmax}_\alpha(z) = \arg\max_{p\in\Delta^m} \left( p^\top z - H_\alpha(p) \right)$$0 and $$\text{entmax}_\alpha(z) = \arg\max_{p\in\Delta^m} \left( p^\top z - H_\alpha(p) \right)$$1 provide $$\text{entmax}_\alpha(z) = \arg\max_{p\in\Delta^m} \left( p^\top z - H_\alpha(p) \right)$$2 mask generation and high recall. Distance-based grouping remains an upper bound on the attainable trade-off, and is computationally efficient when $$\text{entmax}_\alpha(z) = \arg\max_{p\in\Delta^m} \left( p^\top z - H_\alpha(p) \right)$$3.

In both MT and MLM tasks, Sparsefinder’s distance and clustering variants define new Pareto fronts in sparsity–recall and task performance, outperforming established fixed and learned pattern baselines. The evidence indicates that efficient prediction of content-driven sparse patterns is not only feasible but more accurate than prior fixed or dynamically-parameterized sparse frameworks.

A plausible implication is the value of reporting sparsity–recall curves in future comprehensive sparse attention benchmarks, and motivating further research on hardware-sensitive implementations for these predicted mask patterns (Treviso et al., 2021).

6. Implications and Position within Sparse Attention Research

Sparsefinder proposes a new paradigm for attention efficiency: instead of imposing fixed or random sparsity patterns (as in Longformer, BigBird, Routing Transformer, Reformer), it learns to predict the specific nonzeros of entmax attention from low-dimensional projections and fast grouping. The method fully retains the content sensitivity of the original model and permits exact recovery of all truly attended positions.

This approach advances the field beyond both prior fixed-sparsity (e.g., local windows, global tokens) and recent learned-sparsity methods (e.g., Routing Transformer), presenting a content-driven, highly efficient path for sparse attention computation. The separation of learning projections (one-time, lightweight) and inference-time mask prediction (highly parallelizable, hardware-friendly) positions Sparsefinder as an efficient component for large-scale transformer inference and training when entmax or sparsemax attention is used.

The analysis and empirical results of Treviso et al. suggest that further research should both incorporate such mask prediction benchmarks and explore hardware/software co-designs to optimally leverage Sparsefinder-style masking in high-throughput transformer systems (Treviso et al., 2021).