[Paper Review] Tunable Efficient Unitary Neural Networks (EUNN) and their application to RNNs

Jing Li, Yichen Shen+6 more|arXiv (Cornell University)|Dec 15, 2016

Speech Recognition and Synthesis22 references97 citations

TL;DR

This paper introduces Tunable Efficient Unitary Neural Networks (EUNNs), a novel RNN architecture that parametrizes the full unitary matrix space with $Ó(1)$ computational cost per parameter, enabling stable training without gradient vanishing/exploding. EUNNs achieve state-of-the-art performance on long-sequence tasks like the copying task, permuted MNIST, and TIMIT speech prediction, outperforming both LSTMs and prior unitary RNNs in accuracy and training speed.

ABSTRACT

Using unitary (instead of general) matrices in artificial neural networks (ANNs) is a promising way to solve the gradient explosion/vanishing problem, as well as to enable ANNs to learn long-term correlations in the data. This approach appears particularly promising for Recurrent Neural Networks (RNNs). In this work, we present a new architecture for implementing an Efficient Unitary Neural Network (EUNNs); its main advantages can be summarized as follows. Firstly, the representation capacity of the unitary space in an EUNN is fully tunable, ranging from a subspace of SU(N) to the entire unitary space. Secondly, the computational complexity for training an EUNN is merely $\mathcal{O}(1)$ per parameter. Finally, we test the performance of EUNNs on the standard copying task, the pixel-permuted MNIST digit recognition benchmark as well as the Speech Prediction Test (TIMIT). We find that our architecture significantly outperforms both other state-of-the-art unitary RNNs and the LSTM architecture, in terms of the final performance and/or the wall-clock training speed. EUNNs are thus promising alternatives to RNNs and LSTMs for a wide variety of applications.

Motivation & Objective

• Address the vanishing and exploding gradient problems in RNNs by leveraging unitary weight matrices, which preserve gradient stability during backpropagation.
• Overcome limitations of prior unitary RNNs that rely on projection-based optimization or restricted subspaces, which are computationally expensive or less expressive.
• Develop a parametrization of the full unitary group that allows tunable capacity—from subspace to full unitary space—while maintaining computational efficiency.
• Demonstrate superior performance on long-term sequential learning benchmarks, including the copying task, pixel-permuted MNIST, and speech prediction on TIMIT, with reduced parameter counts and faster training.

Proposed method

• Propose a new parametrization of unitary matrices using a product of rotation matrices with tunable structure, enabling full coverage of the unitary space $U(N)$ or its subspaces.
• Design the architecture such that each parameter update and gradient computation requires only $Ó(1)$ operations per parameter, achieving linear-time training complexity.
• Implement two variants: a tunable-style EUNN for flexible capacity control and an FFT-like style for minimal parameterization and efficient approximation.
• Use a reparameterization strategy that avoids projection steps during training, eliminating the need for iterative unitary constraints and reducing computational overhead.
• Apply the EUNN to recurrent networks by replacing the standard hidden-to-hidden weight matrix with a unitary matrix parametrized via rotation matrices.
• Train using standard backpropagation with gradient computation via the chain rule, where the Jacobian of the unitary transformation is computed efficiently using closed-form derivatives.

Experimental results

Research questions

• RQ1 Can a unitary RNN architecture be designed to span the full unitary group $U(N)$ without relying on projection-based optimization?
• RQ2 Does a unitary RNN with $Ó(1)$ computational cost per parameter achieve better training efficiency and performance than existing methods on long-sequence tasks?
• RQ3 Can the representation capacity of a unitary RNN be tuned to match task-specific needs, such as low capacity for short-term memory and high capacity for long-term dependencies?
• RQ4 How does the EUNN architecture compare to LSTMs and other unitary RNNs in terms of final accuracy and wall-clock training speed on real-world benchmarks like TIMIT and permuted MNIST?
• RQ5 Is the EUNN architecture robust to different matrix ordering schemes in the product decomposition, while maintaining computational efficiency and full coverage of the unitary space?

Key findings

• The EUNN architecture achieves $Ó(1)$ computational cost per parameter for both forward and backward passes, making it $Ó(\log N)$ more efficient than the prior full-space unitary RNN method (Wisdom et al., 2016).
• On the copying task, the EUNN reaches 100% accuracy with a hidden size of 128, outperforming both LSTM and prior unitary RNNs.
• On pixel-permuted MNIST, the EUNN achieves 97.5% test accuracy with only 33k parameters, significantly outperforming the LSTM baseline and prior unitary RNNs.
• On the TIMIT speech prediction task, the full-capacity EUNN achieves a test MSE of 51.9, outperforming the LSTM (test MSE 54.5) and other unitary RNN variants.
• The EUNN with tunable capacity (e.g., 128(2) or 128(32)) shows that smaller subspaces are sufficient for short-term memory tasks like copying, while larger subspaces are needed for complex tasks like speech prediction.
• The FFT-style EUNN variant achieves comparable performance to the full EUNN with fewer parameters, demonstrating that efficient approximation is viable without sacrificing accuracy.