FlashSketch: GPU Efficiency & Interactive FG-SBIR

• FlashSketch is a dual-approach framework that integrates GPU-optimized sparse sketching with RL-driven sketch-based image retrieval to accelerate interactive sketching tasks.
• It employs the BlockPerm-SJLT method with a novel CUDA kernel to improve memory access patterns and achieve a global geometric mean speedup of approximately 1.7× over traditional GPU SJLT implementations.
• The RL component uses CNN+attention and fast ANN search to optimize early retrieval, delivering sub-50ms response times in fine-grained sketch-based image retrieval applications.

FlashSketch encompasses two independent advances, each targeting "sketching" but within disparate technical domains: (1) randomized numerical linear algebra (RandNLA) for efficient high-dimensional linear mappings on GPUs (Dwaraknath et al., 2 Feb 2026), and (2) reinforcement learning-driven, real-time fine-grained sketch-based image retrieval (FG-SBIR) (Bhunia et al., 2020). Each line of work has yielded a flagship system named “FlashSketch,” and both share the motivation of accelerating interactive or computational sketching primitives by algorithm-implementation co-design.

1. High-Performance Sparse Sketching on GPUs: BlockPerm-SJLT and FlashSketch Kernels

The first FlashSketch system targets the bottleneck of irregular arithmetic and memory access in GPU implementations of sparse Johnson–Lindenstrauss transforms (SJLTs). Classical SJLTs leverage random sparsity to reduce the cost of sketching $A\in\mathbb{R}^{d\times n}$ via $S\in\mathbb{R}^{k\times d}$, but this random sparsity leads to highly irregular global memory accesses and update patterns on GPUs, negating arithmetic savings through poor bandwidth utilization and excessive use of global atomic operations (Dwaraknath et al., 2 Feb 2026).

FlashSketch introduces a co-designed strategy:

• Algorithmic innovation: A new sparse sketch family, BlockPerm-SJLT, whose block-permutation structure allows hardware-efficient memory access.
• Kernel implementation: An optimized CUDA kernel that exploits the sketch structure to maximize memory bandwidth and avoid expensive global atomics.

This design yields a tunable trade-off—exposing a parameter $\kappa$—between sketching robustness and GPU efficiency, and achieves global geometric mean speedups of approximately 1.7× over prior GPU SJLT implementations across standard RandNLA and end-to-end machine learning data attribution tasks (Dwaraknath et al., 2 Feb 2026).

2. BlockPerm-SJLT Construction and Theoretical Guarantees

The BlockPerm-SJLT extends classical SJLTs by partitioning both input and output coordinates into $M$ blocks. Each output block $g$ aggregates input from a neighborhood $\mathcal{N}(g)$ consisting of $\kappa$ input blocks, which are determined by $\kappa$ edge-disjoint permutations of the $M$ blocks. For each pair $(g, h)$ with $h \in \mathcal{N}(g)$, an independent intra-block SJLT $\Phi_{g, h}$ assigns $s$ nonzeros per column.

Key properties:

• Each input block participates in $\kappa$ output blocks, creating a $\kappa$-regular bipartite sketch graph.
• The sketch matrix $S$ possesses block locality and structure, with $S_{g, h} = (1/\sqrt{\kappa})\Phi_{g, h}$ if $h\in \mathcal{N}(g)$, ensuring $\kappa s$ nonzeros per column, each of magnitude $1/\sqrt{\kappa s}$.

The principal theoretical result establishes an oblivious subspace embedding (OSE) guarantee. "Neighborhood coherence" $\mu_{\mathrm{nbr}}(U;\pi)$, a generalization of block coherence, quantifies the localization of orthonormal subspaces with respect to block neighborhoods. The OSE theorem asserts that, for $U\in\mathbb{R}^{d\times r}$:

• If $$k \gtrsim (\mu_{\mathrm{nbr}}(U;\pi)/\epsilon^2)(r+\log(1/\delta))$$ and $\kappa s \gtrsim (1/\epsilon)(r+\log(1/\delta))$, then $$\left\|U^\top S^\top S U - I_r\right\|_2 \leq \epsilon$$ with probability at least $1-\delta$.
• For highly mixed $U$, $\mu_{\mathrm{nbr}}(U;\pi)\approx1$, recovering the classical dimension bounds for OSEs (Dwaraknath et al., 2 Feb 2026).

3. FlashSketch CUDA Kernel Architecture and GPU Optimization

FlashSketch’s CUDA kernel is organized around block-wise computation. Noteworthy features include:

• Data is laid out in contiguous blocks, reducing memory scattering.
• The sparsity pattern (permutations, hash indices, and signs) is computed on-the-fly using integer arithmetic, eliminating storage and global memory lookups for $S$.
• The kernel uses a 2D grid over output blocks and column tiles, processes each output block via a shared-memory tile, and performs all atomic additions within thread-block shared memory.
• Once all relevant input blocks are fully processed, the shared-memory output tile is normalized and committed to global memory via a coalesced write.

Critical optimizations:

• No global atomics are required for the vast majority of operations.
• Register tiling and vector loads amortize address calculations.
• On-the-fly randomness with no global lookups further reduces cache pressure.

Block neighborhood generation uses an efficient affine-map permutation with full period, multiplying performance—though these are not fully independent, empirical observations indicate negligible practical difference for $\kappa \ll M$ (Dwaraknath et al., 2 Feb 2026).

4. Parameter Trade-offs, Empirical Performance, and Use Cases

Parameter $\kappa$ controls the degree of mixing versus bandwidth utilization:

• Increasing $\kappa$ improves subspace mixing and "smooths" $\mu_{\mathrm{nbr}}$, reducing the number of rows $k$ needed for a target error, at the cost of greater read and shared-atomic workload.
• $\kappa=1$ delivers maximum bandwidth but suffers from high $\mu_{\mathrm{nbr}}$ (worse error).
• $\kappa$ in the range 4–8 achieves $\mu_{\mathrm{nbr}}\approx1$ with modest bandwidth drop (to 60–70% of peak).

Empirical benchmarks on RTX 4090 demonstrate:

• Gram approximation for GPT2-medium ($d=16384, n=1024, k\approx2048$): GraSS kernel takes 2.1 ms for error 0.05; FlashSketch ($\kappa=4$) achieves 1.3 ms at the same error.
• Aggregated across tasks—Gram approximation, OSE spectral-error test, sketch-and-ridge, and sketch-and-solve—FlashSketch achieves a global geomean speedup of approximately 1.7× over the next-best baseline.
• In end-to-end ML pipelines (e.g., MNIST MLP feature-caching), FlashSketch achieves up to 3.2× per-sample speedup with no degradation in data attribution metrics (Dwaraknath et al., 2 Feb 2026).

Typical parameters: $\kappa$ between 2–8, $s$ between 1–4, with per-shape tuning of block and tile sizes. For small $k$, occupancy drops and a split-$B_c$ fallback using global atomics is required.

Use cases include high-dimensional least squares and ridge regression with $d\gg n$ and end-to-end ML pipelines requiring repeated sketching.

5. FlashSketch for On-the-Fly Fine-Grained Sketch-Based Image Retrieval

A second, independent system named FlashSketch addresses latency and usability in FG-SBIR (Bhunia et al., 2020). Here, the challenge is to retrieve the target image with minimal sketch input and deliver top-$q$ results interactively as the user draws.

Central mechanisms:

• Reinforcement learning (RL): The FlashSketch RL framework observes the partial rasterized sketch $s_t$ at drawing step $t$, encodes this via a CNN+attention+actor (inception-V3 backbone), and outputs an embedding $a_t$.
• Policy and reward: The policy $\pi_\Theta(a_t|s_t)$ is Gaussian, parameterized by a fully-connected actor head. RL training (via PPO) directly optimizes the expected cumulative reward—high early retrieval rank (local reward) and stable ranking consistency (global reward via Kendall’s $\tau$ penalty).
• Retrieval mechanism: At each stroke batch, the system computes $a_t$, finds nearest neighbors in a pre-computed gallery using approximate nearest-neighbor (Faiss IVFADC) over $\mathbb{R}^{64}$, and presents the top-$q$ images in $<10$ ms.

Experimental results (on QMUL‐Chair‐V2 and QMUL‐Shoe‐V2) demonstrate that RL-based FlashSketch outperforms various baselines both in final accuracy (Acc@5/10) and early retrieval metrics (mean area under curve for ranking percentile and inverse rank), with efficacy particularly pronounced in the initial 30–50% of strokes (Bhunia et al., 2020).

6. System Design, Latency Considerations, and Practical Implications

The real-time FlashSketch system for FG-SBIR is composed of:

1. Drawing client that batches vector strokes.
2. Rasterizer for input resizing.
3. CNN+attention+actor feature encoder optimized for incremental processing.
4. Fast ANN search over gallery photo embeddings.
5. Responsive UI with uncertainty visualization and early stop logic.
6. Async pipelining: while the user draws, inference and search proceed for the previous batch.

Critical design points:

• End-to-end latency is constrained to $\leq50$ ms per stroke to preserve interactive responsiveness.
• The system employs GPUs or SIMD-enabled CPUs for both CNN inference and ANN search.
• Scalability for massive galleries is achieved via coarse-to-fine hashing.

A plausible implication is that the integration of RL reward shaping and fast feature-ANN search can bring sketch-based search to parity with text/tag queries in user-perceived latency, substantially enhancing usability for practical FG-SBIR applications.

---

In summary, FlashSketch—in both its RandNLA/GPU and RL-FG-SBIR manifestations—implements a co-design philosophy: structuring algorithms to maximize hardware efficiencies and/or optimize early-stage output quality. In sparse sketching for GPUs, this enables OSE-theoretic guarantees at unprecedented performance, while in image retrieval, FlashSketch facilitates minimal-interaction, high-accuracy search through RL-driven embedding evolution and rapid nearest-neighbor querying (Dwaraknath et al., 2 Feb 2026, Bhunia et al., 2020).