PSSA: A non-transformer language model written from scratch in Rust (github.com)

PSSA, a newly developed language model implemented entirely in Rust, diverges from traditional transformer architectures by utilizing a recurrent state-space mechanism paired with an episodic memory bank. Unlike transformers that compute the relationship between every token, causing exponential growth in computational demand with increasing sequence lengths, PSSA processes tokens in a linear fashion and can retrieve data from its memory rather than re-reading contexts. This innovative approach has allowed PSSA to outperform transformers in both training speed and efficiency, achieving a lower training cross-entropy of 3.98 compared to a transformer model’s 4.43 over 12.7 million tokens from WikiText-103. The significance of PSSA lies not only in its faster learning curve and text generation—achieving token generation speeds twelve times faster on the same CPU—but also in its potential contributions to the AI/ML community. By exploring a fundamentally different architecture, PSSA could inspire new models that prioritize efficiency and generalization over sheer power, particularly relevant in resource-constrained environments. Additionally, its open-source foundations invite collaboration, allowing researchers to explore further optimizations, experiment with memory bank functionalities, and validate its performance across larger datasets.