Work

Rekurrente Architekturen und hardware-effiziente Implementierungen für Deep Learning auf großen Skalen

Year: 2025
Type: dissertation
Abstract: Transformers are the de facto standard architecture in most current Machine Learning systems, from Natural Language Processing, to Computer Vision and Physics Simulation. Their benefit of being able to access all the past or other context explicitly is also their biggest downside: Their compute scales quadratically in the context length. Recurrent architectures had been established before the advent of Transformers, but were inferior in capabilities for a long time. This work extends ideas from the dominant, recurrent architecture before the Transformer, Long Short-Term Memory (LSTM), towards architectures that are recurrent and matching or surpassing Transformers' capabilities. In xLSTM, the gating structure of LSTMs is modified for enhanced updating capabilities. The superiority of non-linear RNNs, like the proposed sLSTM, is shown for logic and formal language tasks that need state tracking capabilities. Additionally, xLSTM introduces a variant, mLSTM, with extended (associative) memory, that has improved associative recall performance and allows xLSTM to outperform Transformers and other recurrent alternatives on language modeling. The FlashRNN library provides an efficient implementation in CUDA of traditional non-linear RNNs like sLSTM on modern hardware, analyzing the I/O bottlenecks, re-using recurrent matrices stored at the register level, and leveraging matrix multiplication acceleration for large speed-ups. In xLSTM-7B, such an xLSTM model is scaled to the seven billion parameter scale trained on 2.3T tokens. For this, the backbone architecture around the mLSTM is optimized for speed, leading to the fastest dense model at this scale, with additional analyses on long-context training and capabilities. Finally, in the pLSTM work, the notion of linear RNNs is extended to multi-dimensional RNNs, leading to parallelizable Linear Source Transition Mark networks that can operate on general DAGs. It shows how moving the gates from in between nodes to in between edges of a graph enables parallelization that can be even in the form of generalized matrix multiplications in the case of grid structures, e.g. images. On top, it is the first to analyze long-range stability and vanishing/exploding activations/gradients of linear networks on DAG structures, leading to separate stabilization modes, that cover directional propagation and diffusive distribution. These works altogether provide recurrent, and therefore linear alternatives to currently dominant quadratic architectures, with the downside of a compressed history or context integration. Hardware-efficient implementations enable models to be scaled to sizes at commercial levels (7B) and speeds superior to existing alternatives. Beyond their higher expressivity, their linear scaling should ultimately lead to faster deep learning models that need less memory, compute, and energy.
Author: Pöppel, Korbinian
Language: English
Open Access status: green
OpenAlex ID: W7111935471
Cites:
Cited by:
Related to: 10