Randomized Sketching is Robust to Low-Precision Rounding on GPUs

2026-06-18

AI Summary

This work investigates the implementation of Sparse Unaware Subspace Embeddings (SparseStack) under FP16 mixed-precision arithmetic on GPUs, where performance is typically constrained by memory bandwidth and atomic operations, and reduced precision may degrade embedding quality. The study systematically evaluates the impact of deterministic rounding, stochastic rounding, and dithered quantization on subspace distortion and least-squares accuracy. Empirical results demonstrate that all FP16 SparseStack variants preserve nearly identical embedding fidelity across coherent, incoherent, and adversarial problem instances, indicating that the sketching distribution—not the rounding strategy—primarily governs numerical accuracy. Among the approaches examined, deterministic rounding incurs the lowest computational overhead, establishing it as the optimal choice. These findings validate an efficient and practical pathway for deploying FP16 SparseStack in large-scale applications.

Abstract

Randomized sketching is a core primitive in randomized numerical linear algebra. On modern hardware architectures, in particular on GPUs, the performance of sparse sketches is limited by memory traffic and atomic accumulation rather than floating-point throughput. This makes sketching a natural target for mixed precision, provided that low-precision accumulation does not degrade the embedding quality. We study mixed-precision GPU implementations of sparse oblivious subspace embeddings, focusing on a SparseStack generalization of the GPU CountSketch kernel of Higgins et al. SparseStack improves embedding quality relative to CountSketch on coherent inputs, but its additional nonzeros per column increase atomic-update contention and reduce throughput. We therefore implement FP16 SparseStack variants using deterministic round-to-nearest, exact stochastic rounding, and dithered rounding, and compare them with FP32 SparseStack, CountSketch, mixed-precision CountSketch, and FlashSketch. Our main empirical finding is that, for the tested regimes, SparseStack embedding quality is insensitive to the FP16 rounding rule. Deterministic, stochastic, and dithered rounding FP16 SparseStack produce nearly identical subspace distortion and sketch-and-solve least-squares accuracy across incoherent, coherent, and adversarial test problems. The dominant accuracy factor is the sketch distribution rather than the quantization rule: SparseStack variants substantially improve distortion on coherent inputs, while all methods behave similarly on incoherent inputs. Since deterministic rounding has the lowest overhead, it provides the best performance--accuracy tradeoff among the FP16 SparseStack variants.

Problem

• randomized sketching
• low-precision rounding
• GPU
• subspace embedding
• mixed precision

Innovation

• mixed-precision
• randomized sketching
• SparseStack
• GPU acceleration
• low-precision rounding