Fast Heavy Inner Product Identification Between Weights and Inputs in Neural Network Training

In this paper, we consider a heavy inner product identification problem, which generalizes the Light Bulb problem~(\cite{prr89}): Given two sets $A\subset \{-1,+1{\}}^d$ and $B\subset \{-1,+1{\}}^d$ with $\mathrm{\mid }A\mathrm{\mid }=\mathrm{\mid }B\mathrm{\mid }=n$, if there are exact $k$ pairs whose inner product passes a certain threshold, i.e., $\{(a_1,b_1),\cdots ,(a_k,b_k)\}\subset A\times B$ such that $\mathrm{\forall }i\in [k],⟨a_i,b_i⟩\ge \rho \cdot d$, for a threshold $\rho \in (0,1)$, the goal is to identify those $k$ heavy inner products. We provide an algorithm that runs in $O(n^{2\omega \mathrm{/}3+o(1)})$ time to find the $k$ inner product pairs that surpass $\rho \cdot d$ threshold with high probability, where $\omega$ is the current matrix multiplication exponent. By solving this problem, our method speed up the training of neural networks with ReLU activation function.