Block-Max Pruning: Model & Retrieval Efficiency

• Block-Max Pruning (BMP) is a structured technique that partitions data into blocks and retains the most salient values for compression and indexing.
• It creates regular sparsity patterns that reduce index storage overhead and enable parallel hardware decoding.
• Empirical results show BMP delivers significant speedups in neural compression and learned retrieval with minimal accuracy loss.

Block-Max Pruning (BMP) is a structured pruning and indexing technique applied in two major domains: neural network model compression and fast search in learned sparse retrieval systems. In both contexts, BMP is predicated on the use of blockwise upper bounds—whether on tensor weights or document scores—to tightly balance computational efficiency, regularity, and model/task performance. The core idea is to partition a structure (weights or document IDs) into small blocks and, for each, retain or reason about only the most salient value, enabling aggressive compression or rapid pruning while retaining accuracy or retrieval quality.

1. Motivation and Structural Insights

In neural networks, irregular (unstructured) pruning achieves high compression by retaining weights with the largest magnitudes regardless of their position, but produces random sparsity patterns. This irregularity results in significant index-storage overhead (e.g., using CSR encoding) and index-decoding cost (requiring complex hardware for arbitrary lookup), limiting throughput and efficiency on modern devices. Conversely, traditional structured pruning improves regularity but often sacrifices the pruning ratio, as only predetermined blocks or patterns of weights survive. Detailed analysis of irregularly pruned neural weight matrices reveals heterogeneous row densities (up to 14× variation), heightened sensitivity of dense rows to pruning, and positional patterns where many zero-weight blocks are nearly as significant as those with retained weights. These observations motivate the adoption of a local "block-max" selection scheme that preserves most salient values while reinstating regularity.

In the search and retrieval domain, traditional dynamic pruning strategies (such as WAND, MaxScore, and Block-Max WAND) were designed for short queries and smooth score distributions typical of term-weighted models (e.g., BM25). Learned sparse retrieval—e.g., using SPLADE, ESPLADE, or uniCOIL—produces much longer queries (50–200 tokens), smaller vocabularies, and quantized/skewed score distributions, causing traditional pruning to either process excessive postings or to over-prune, substantially reducing efficiency or accuracy. BMP in this context introduces block filtering and upper bounds over blockwise document ranges, enabling safe or approximate early termination and restoring efficiency while maintaining exact top-k guarantees (Mallia et al., 2024).

2. Methodologies: Block-Max Weight Masking and Block-Range Pruning

Block-Max Weight Masking (BMWM) in Neural Network Pruning

Given a weight matrix $W \in \mathbb{R}^{R \times C}$, each row is partitioned into $B = \lfloor C/m \rfloor$ contiguous blocks of width $m$. Let $\mathcal{B}_{i,b}$ be the $b$th block of row $i$. The binary mask $M$ selects exactly one weight per block: $$M_{i,j} = \begin{cases} 1 & \text{if } (i,j) = \arg\max_{(p,q) \in \mathcal{B}_{i,b}} |W_{p,q}| \ 0 & \text{otherwise} \end{cases}$$ This ensures one nonzero per block, yielding a per-row density of $1/m$ and global sparsity of $1 - 1/m$. Smaller $B = \lfloor C/m \rfloor$0 increases pruning ratio but can reduce the likelihood of retaining salient weights within blocks. The time complexity is $B = \lfloor C/m \rfloor$1, and decoding each block reduces to a fixed-offset multiplexer.

Density-Adaptive Regular-Block Pruning (DARB)

To accommodate row sensitivity and density heterogeneity, DARB first irregularly prunes $B = \lfloor C/m \rfloor$2 to a global target density $B = \lfloor C/m \rfloor$3. For each row $B = \lfloor C/m \rfloor$4, the irregularly derived density $B = \lfloor C/m \rfloor$5 is rounded to the nearest power-of-two fraction $B = \lfloor C/m \rfloor$6. If $B = \lfloor C/m \rfloor$7, round up to preserve sensitivity; if $B = \lfloor C/m \rfloor$8, round down. Set the row’s block size $B = \lfloor C/m \rfloor$9 (constrained to $m$0), then apply BMWM with per-row $m$1 (Ren et al., 2019).

Block-Max Pruning in Sparse Retrieval

Given $m$2 documents (docIDs $m$3), choose block size $m$4 and partition into $m$5 blocks $m$6. For each vocabulary term $m$7, construct a block-max impact array $m$8 (with $m$9 if $\mathcal{B}_{i,b}$0 is absent). For a query $\mathcal{B}_{i,b}$1 with term-weights $\mathcal{B}_{i,b}$2, define a block upper bound: $\mathcal{B}_{i,b}$3 Block processing is ordered by descending $\mathcal{B}_{i,b}$4, and early safe termination occurs once the threshold $\mathcal{B}_{i,b}$5 (from the current top-$\mathcal{B}_{i,b}$6 min-heap) is no less than the maximum remaining unprocessed block upper bound. An approximate mode with parameter $\mathcal{B}_{i,b}$7 allows early stop when $\mathcal{B}_{i,b}$8 (Mallia et al., 2024).

3. Theoretical Guarantees and Trade-offs

BMP for neural pruning with block-size $\mathcal{B}_{i,b}$9 guarantees exactly one nonzero per block, hence a pruning ratio of $b$0. DARB further tailors sparsity, limiting sensitivity-induced degradation, and keeps encoding efficient by restricting block sizes to powers of two. The encoding overhead is modest: per-weight index storage is $b$1 bits, with total index storage for models such as 1500x1500 LSTM being 0.63 MB, well below CSR overheads.

In retrieval, the dominant cost is $b$2 for block upper bounds and $b$3 for partial block sort. In the worst case, retrieval degenerates to full scan ($b$4), but in common practice, rapid upper bound decay and threshold increases yield visits to only $b$5–$b$6 of blocks, leading to substantial speedups. The termination guarantee is strict—no document not examined can enter the top-$b$7 result set if the safe stopping criterion is fulfilled (Mallia et al., 2024).

4. Empirical Performance

Neural Network Model Compression

DARB achieves pruning ratios of $b$8–$b$9—substantially higher than earlier structured methods—while delivering minimal accuracy loss. For LSTM (PTB), DARB-1/2 achieves $i$0 pruning with test PPL increases of $i$1; on AlexNet-FC (ImageNet), DARB-I/II reaches $i$2 with at most $i$3 Top-5 accuracy drop, approaching irregular pruning (24×, no loss). In decoding efficiency, DARB’s 1,478 activations/cycle (13.14× prune) far exceeds block $i$4 (7.48×, 103 activates/cycle) and $i$5 (5.37×, 413 acts/cycle)—a $i$6 and $i$7 improvement, respectively (Ren et al., 2019).

Learned Sparse Retrieval

On MS MARCO with SPLADE, ESPLADE, and uniCOIL models, BMP (block size $i$8) achieves mean response times of $i$9, $M$0, and $M$1 ms for exact top-10 retrieval, compared to MaxScore (120.6/13.6/14.5 ms), BMW (614.2/10.8/12.4 ms), IOQP (79.1/27.3/34.8 ms), and Anytime (80.6/8.2/8.3 ms), representing up to $M$2 speedup over BMW and at least $M$3 improvement over all baselines at $M$4. Approximate retrieval with BMP ($M$5, $M$6) achieves RR@10 of $M$7 with $M$8 ms response time, close to exhaustive search and surpassing IOQP (RR@10 $M$9 at $$M_{i,j} = \begin{cases} 1 & \text{if } (i,j) = \arg\max_{(p,q) \in \mathcal{B}_{i,b}} |W_{p,q}| \ 0 & \text{otherwise} \end{cases}$$0 ms), Anytime (RR@10 $$M_{i,j} = \begin{cases} 1 & \text{if } (i,j) = \arg\max_{(p,q) \in \mathcal{B}_{i,b}} |W_{p,q}| \ 0 & \text{otherwise} \end{cases}$$1 at $$M_{i,j} = \begin{cases} 1 & \text{if } (i,j) = \arg\max_{(p,q) \in \mathcal{B}_{i,b}} |W_{p,q}| \ 0 & \text{otherwise} \end{cases}$$2 ms) (Mallia et al., 2024).

5. Practical Implications and Deployment

In neural pruning, BMP reduces the index storage per retained weight to $$M_{i,j} = \begin{cases} 1 & \text{if } (i,j) = \arg\max_{(p,q) \in \mathcal{B}_{i,b}} |W_{p,q}| \ 0 & \text{otherwise} \end{cases}$$3 bits, typically $$M_{i,j} = \begin{cases} 1 & \text{if } (i,j) = \arg\max_{(p,q) \in \mathcal{B}_{i,b}} |W_{p,q}| \ 0 & \text{otherwise} \end{cases}$$4 bits, with minimal overhead over rigid block pruning and orders-of-magnitude less than irregular (CSR-based) encoding. Hardware decoding becomes parallelizable, leveraging small, fixed-width multiplexers. DARB’s guidelines: conduct initial irregular pruning, measure row densities, set target global density, round per-row densities to powers of two, assign block sizes accordingly, apply BMWM, and fine-tune as needed.

In retrieval, BMP requires constructing a block-max index for each term, incurring a storage overhead of $$M_{i,j} = \begin{cases} 1 & \text{if } (i,j) = \arg\max_{(p,q) \in \mathcal{B}_{i,b}} |W_{p,q}| \ 0 & \text{otherwise} \end{cases}$$5–$$M_{i,j} = \begin{cases} 1 & \text{if } (i,j) = \arg\max_{(p,q) \in \mathcal{B}_{i,b}} |W_{p,q}| \ 0 & \text{otherwise} \end{cases}$$6 GB over the forward index, depending on $$M_{i,j} = \begin{cases} 1 & \text{if } (i,j) = \arg\max_{(p,q) \in \mathcal{B}_{i,b}} |W_{p,q}| \ 0 & \text{otherwise} \end{cases}$$7. Query-time computation and block filtering are amenable to SIMD and are substantially more cache efficient. The block size $$M_{i,j} = \begin{cases} 1 & \text{if } (i,j) = \arg\max_{(p,q) \in \mathcal{B}_{i,b}} |W_{p,q}| \ 0 & \text{otherwise} \end{cases}$$8 and early-stop parameter $$M_{i,j} = \begin{cases} 1 & \text{if } (i,j) = \arg\max_{(p,q) \in \mathcal{B}_{i,b}} |W_{p,q}| \ 0 & \text{otherwise} \end{cases}$$9 must be tuned for specific workloads and models. Hybrid block forward indexing further reduces pointer chasing overhead, crucial for long postings in learned sparse settings.

6. Comparisons, Limitations, and Forward Directions

Traditional structured block pruning (for neural networks) constrained pruning flexibility and failed to capture the nuanced density sensitivity revealed by irregular sparsity. DARB, by adopting density-adaptive block sizes, matches irregular methods’ pruning ratios while ensuring hardware-friendly patterns. Similarly, for search, legacy dynamic pruning schemes perform poorly on learned sparse indexes: MaxScore is overwhelmed by query expansion, and BMW incurs heavy overhead from many small blocks. BMP, by aggressively filtering $1/m$0–$1/m$1 of document ranges before detailed evaluation and supporting both exact and approximate retrieval, outperforms all tested baselines in both runtime and top-$1/m$2 precision.

Limitations include additional index storage in both domains and the necessity to tune block sizes and thresholds for workload-specific performance. The retrieval BMP implementation is currently single-threaded; parallel block evaluation and GPU/SIMD acceleration are promising future directions. Potential comparisons with graph-based ANN for sparse vectors and further automation of parameter selection constitute natural next steps (Ren et al., 2019, Mallia et al., 2024).