Chaudhary, Sanyam
Axiom is a recurrent neural network that learns a unitary transition matrix via a product of Householder reflections, perfectly preserving long-term dependencies without a forget gate.
Traditional RNNs (LSTM, GRU) struggle to learn long-range dependencies over thousands of steps due to vanishing/exploding gradients. The forget gate allows selective forgetting, but in some tasks, information loss is critical.
The hidden-to-hidden transition matrix is parameterized as a product of k Householder reflections, maintaining a strict unitary matrix. A closed-form parallel forward pass is derived, reducing the unitary recurrence to a cumulative sum in a rotated eigenbasis. A simplified single-sided rotation is introduced for XLA optimization.
On the delayed copy task (T=1000), Axiom achieves 76.5-99.9% accuracy with 8,584 parameters (LSTM: 111,368 parameters, 12.5-13.5%). On the Adding Problem, MSE 0.00046 (LSTM 0.00214). When attached to GPT-2, it retrieves facts from 7 chunks back with 62.3% accuracy (baseline 8.3%). The boundary where LSTM's forget gate is advantageous (noise filtering tasks) is experimentally identified.