ZPoles: Fast Fractional Matrix Powers for Optimization via Rational Approximation

Fractional matrix powers underpin gradient preconditioning in Shampoo and momentum transforms in Muon, but their repeated computation can be expensive. Rational approximations use a small number of shifted matrix inverses, whose cost has limited their appeal for GPU training. We propose ZPoles, a method for computing fractional powers with a new mixed-precision linear solver. The solver keeps diagonal operations in higher precision and batches half-precision matrix multiplications across the shifted linear systems. On an RTX 5090 GPU, the solver is 3.7–4.1× faster than cuSOLVER, trading numerical accuracy for runtime. For matrices of size 1024×1024 to 12288×12288, ZPoles computes inverse fourth roots 2.6–8.1× faster than coupled Newton iteration at a fixed 10-iteration budget, with lower error when both use half-precision matrix-multiplication operands. In GPT-2 pretraining, ZPoles reduces observed Shampoo step time relative to coupled Newton, with lower or similar validation loss. Muon and fractional-power Muon provide further evidence of applicability to momentum transforms. These results support rational approximation as a practical computational primitive for GPU-based optimization.