Source-linked AI summary
Tunable Efficient Unitary Neural Networks (EUNN) and their application to RNNs
Li Jing, Yichen Shen, Tena Dubček, John Peurifoy, Scott Skirlo, Yann LeCun, Max Tegmark, Marin Soljačić
TL;DR
RNNs can suffer from vanishing and exploding gradients, limiting long-term information handling, while existing unitary methods face capacity or computational trade-offs. The paper introduces tunable EUNNs that represent unitary spaces efficiently and evaluates them on copying, pixel-permuted MNIST, and TIMIT speech prediction. EUNNs outperform existing RNN architectures on the reported benchmarks and outperform LSTMs on TIMIT speech prediction.
Problem
RNNs face vanishing and exploding gradients, while prior unitary approaches can restrict representation to a subspace or incur costly full-space training.
Method
The paper introduces EUNNs, whose unitary-space representation is tunable from a subspace toward the full space and costs O(1) per parameter.
Results
EUNNs significantly outperform existing RNN architectures on the copying and pixel-permuted MNIST tasks and outperform an LSTM on TIMIT speech prediction.
Takeaways & Limitations
EUNNs demonstrate strong long-term sequential memorization and perform well on real speech-prediction data.
Takeaways & Limitations
Existing full-space unitary training incurs O(N^3) computational cost, while partial-space methods represent only part of the unitary-matrix space.
Abstract
from arXiv · showhide
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.
1. Introduction
RNNs are especially vulnerable to vanishing and exploding gradients because recurrence can create very deep effective networks. The paper introduces EUNNs to represent unitary matrices efficiently and evaluates them across multiple sequential-learning tasks.
- Motivation: RNN recurrence can correspond to thousands or millions of equivalent hidden layers, making gradient instability particularly severe.Eigenvalues far from unity raised to large powers contribute to vanishing or exploding gradients.
- Existing approaches: LSTMs, GRUs, bidirectional RNNs, and gradient clipping address the problem, but none fundamentally solves vanishing and exploding gradients.LSTMs use gates to contain information inside hidden layers, while clipping is often needed to keep gradients reasonable.
- Unitary RNNs: Unitary hidden matrices keep eigenvalue magnitudes at unity, supporting stable propagation through long sequences.Prior work explored unitary initialization and training within tractable unitary subspaces.
- EUNN contribution: EUNN parametrizes the entire unitary-matrix space computationally efficiently while allowing subspace capacity to be tuned by fixing parameters.Existing unitary subspace models appear as special cases of the architecture.
- Efficiency: O(1) computational steps and memory access per parameter let an O(N)-sized hidden layer compute up to the full N × N gradient matrix.This improves on O(N) training complexity for a full-space unitary network and is O(log N) more efficient than a subspace Unitary RNN.
2. Background
An RNN updates hidden states from sequential inputs and uses them to generate outputs, but backpropagation through recurrent transitions can produce vanishing or exploding gradients. The instability depends strongly on the eigenvalues of the hidden-to-hidden matrix.
- Basic RNN architecture: An RNN processes an input sequence by updating its hidden state at each time step, thereby memorizing past information in the hidden layer.The hidden state is updated from the current input and previous hidden state.
- Basic RNN architecture: The hidden preactivation is z(t) = Ux(t) + Wh(t−1), and the hidden state is h(t) = σ(z(t)).Here σ is the nonlinear activation function.
- Gradient propagation: Backpropagation from hidden to hidden layers is the main location of the vanishing or exploding gradient problem.The input-to-hidden and hidden-to-output matrices are comparatively straightforward to train once the recurrent matrix is optimized.
- Gradient propagation: The gradient across time is a product of Jacobian terms D(k)W, so recurrent dynamics accumulate repeatedly during backpropagation.D(k) is the Jacobian of the pointwise nonlinearity.
- Gradient propagation: Eigenvalues λ_i ≫ 1 can cause gradient explosion, whereas λ_i ≪ 1 can cause gradient vanishing when D(k) eigenvalues remain near unity.Either outcome prevents the RNN from working efficiently over large times T.
3. Unitary RNNs
Unitary RNNs constrain recurrent matrices to preserve eigenvalue magnitudes, but existing approaches trade representational capacity against computational cost. The paper motivates EUNN as a way to access broader unitary spaces efficiently.
- Unitary RNNs: Unitary RNNs can address exploding and vanishing gradients by constraining the hidden-to-hidden matrix to a unitary form.The motivation is that unitary matrices preserve eigenvalue magnitudes at unity.
- Factorized constructions: The prior factorized construction combines diagonal phase matrices, reflection matrices, a permutation, and Fourier transforms to produce a unitary recurrent matrix.Because each factor is unitary, their product W is also unitary.
- Partial-space models: O(N) parameters span only part of the O(N^2)-dimensional space of unitary N × N matrices, trading capacity for computational efficiency.Subsequent methods sought to expand the represented space toward O(N^2).
- Full-space models: Full-space training updates W by backpropagation and projects it back onto the unitary space after each update.This approach avoids restricting W to a smaller parameterized subspace.
- Full-space models: O(N^3) computational cost remains a key limitation of full-space unitary training because backpropagation requires N-dimensional matrix multiplication.The cited method is reported as superior on many RNN tasks despite this cost.
4. Efficient Unitary Neural Network (EUNN) Architectures
EUNN parametrizes unitary matrices with tunable capacity, from subspaces to the full unitary space, while enabling efficient RNN computation through structured decompositions and element-wise operations.
- 4.1. Unitary Matrix Parametrization: EUNN represents arbitrary N × N unitary matrices with N^2 parameters, while fixing parameters can restrict the representation to a unitary subspace.The parametrization uses pairwise U(2) rotations and a diagonal matrix.
- 4.2. Tunable space implementation: Choosing the decomposition depth L controls the represented subspace; L = N spans the entire unitary space.Smaller L yields lower-capacity unitary models.
- 4.3. FFT-style approximation: The FFT-style architecture uses log(N) rotation matrices and N log(N)/2 rotational pairs to approximate arbitrary unitary matrices with all input coordinates interacting.Its rotations pair coordinates at increasing distances rather than only adjacent coordinates.
- 4.4. Efficient implementation of rotation matrices: Each rotational layer is implemented using element-wise multiplications and permutations, so computing Wx requires O(NL) operations and memory access.For full-space implementation L = N; for the FFT-style approximation L = log N.
- 4.5. Nonlinearity: The EUNN uses modReLU with a shared trainable bias, and the authors empirically find this nonlinearity performs best.They suggest the bias threshold may filter noise.
5. Experimental tests of our method
The experiments compare EURNN with LSTM and other unitary RNNs on copying, pixel-permuted MNIST, and TIMIT speech prediction. EURNN generally achieves stronger performance, faster learning, or faster wall-clock convergence under the reported settings.
- Experimental setup: The experiments compare EURNN with LSTM, Partial Space URNN, and Projective full-space URNN across synthetic, image, and speech tasks.The evaluated tasks are the Copying Memory Task, pixel-permuted MNIST, and speech prediction on TIMIT.
- Copying Memory Task: For T = 1000, EURNN outperforms LSTM and other unitary RNNs in learnability and convergence rate on the copying task.The comparison uses n = 8, input length M = 10, batch size 128, and RMSProp training.
- Copying Memory Task: Choosing smaller L or the FFT-style EURNN improves convergence efficiency and wall-clock speed over partial and projective unitary methods.The results also report greater robustness and state that full-capacity unitary matrices are unnecessary for this task.
- Pixel-Permuted MNIST Task: On pixel-permuted MNIST, EURNN significantly outperforms LSTM with the same parameter count, learning faster and reaching higher classification accuracy.EURNN also reaches similar accuracy with fewer parameters.
- Speech Prediction on TIMIT dataset: On TIMIT speech prediction, EURNNs generally perform better than LSTMs; full-capacity EURNNs outperform smaller-capacity and FFT-style EURNNs in this task.The task predicts future STFT log-magnitude frames and evaluates models using mean-squared error.
6. Conclusion
EUNN provides efficient, tunable unitary RNNs that outperform existing architectures on sequential-memory benchmarks and LSTM on speech prediction. Its tunability supports different task-specific unitary-space capacities and suggests broader applications to long-term, high-dimensional tasks.
- O(1) per parameter is the EUNN computational cost, O(log N) more efficient than other discussed methods.The method is presented as computationally efficient for training unitary networks.
- EUNN significantly outperforms existing RNN architectures on the Copying and pixel-permuted MNIST tasks with comparable parameter counts.These results demonstrate strong sequential-memory performance over long time periods.
- EUNN outperforms an LSTM on TIMIT speech prediction.The speech task predicts future STFT log-magnitude frames from sampled TIMIT audio.
- Task-specific capacity tuning favors a small unitary subspace for Copying, while MNIST and TIMIT perform better with considerably larger unitary spaces.Different rotation-matrix orderings can span the full unitary space with identical speed and memory performance.
- EUNN may also suit natural-language processing because it efficiently handles long-term correlations and very high dimensionality.