Memorization Capacity of Neural Networks with Conditional Computation
This work addresses a fundamental limit in neural network efficiency for researchers and practitioners, offering theoretical insights into conditional computation's benefits, though it is incremental in building on prior memorization studies.
The paper tackles the problem of memorization capacity in neural networks by analyzing conditional computation, showing that conditional ReLU networks can memorize n input-output relationships using O(log n) operations per input, an almost exponential improvement over the O(sqrt(n)) operations required by unconditional networks, and proves this rate is optimal.
Many empirical studies have demonstrated the performance benefits of conditional computation in neural networks, including reduced inference time and power consumption. We study the fundamental limits of neural conditional computation from the perspective of memorization capacity. For Rectified Linear Unit (ReLU) networks without conditional computation, it is known that memorizing a collection of $n$ input-output relationships can be accomplished via a neural network with $O(\sqrt{n})$ neurons. Calculating the output of this neural network can be accomplished using $O(\sqrt{n})$ elementary arithmetic operations of additions, multiplications and comparisons for each input. Using a conditional ReLU network, we show that the same task can be accomplished using only $O(\log n)$ operations per input. This represents an almost exponential improvement as compared to networks without conditional computation. We also show that the $Θ(\log n)$ rate is the best possible. Our achievability result utilizes a general methodology to synthesize a conditional network out of an unconditional network in a computationally-efficient manner, bridging the gap between unconditional and conditional architectures.