Improving the Robustness of Neural Multiplication Units with Reversible Stochasticity
This addresses a robustness issue in specialized neural modules for arithmetic, which is incremental as it builds on existing NMUs to enhance their reliability.
The paper tackled the problem of Neural Multiplication Units (NMUs) failing to learn simple multiplication tasks reliably across different training ranges, and proposed a stochastic NMU (sNMU) that uses reversible stochasticity to avoid undesirable optima and converge to true solutions, improving robustness.
Multilayer Perceptrons struggle to learn certain simple arithmetic tasks. Specialist neural modules for arithmetic can outperform classical architectures with gains in extrapolation, interpretability and convergence speeds, but are highly sensitive to the training range. In this paper, we show that Neural Multiplication Units (NMUs) are unable to reliably learn tasks as simple as multiplying two inputs when given different training ranges. Causes of failure are linked to inductive and input biases which encourage convergence to solutions in undesirable optima. A solution, the stochastic NMU (sNMU), is proposed to apply reversible stochasticity, encouraging avoidance of such optima whilst converging to the true solution. Empirically, we show that stochasticity provides improved robustness with the potential to improve learned representations of upstream networks for numerical and image tasks.