SUBLLM: A Novel Efficient Architecture with Token Sequence Subsampling for LLM

TL;DR

SUBLLM introduces a decoder-only LLM architecture that dynamically allocates compute via learnable token subsampling, upsampling, and a bypass connection. By shortening the input sequence with a token-importance scorer and reconstructing it with differentiable scaling, SUBLLM achieves substantial training and inference speed-ups while reducing GPU memory, without sacrificing few-shot performance. The bypass module further accelerates convergence, and experiments show up to roughly 26% training speed-up and 37% inference speed-up, with larger gains as context windows grow; results generalize across 0.25B and 1.3B models on SlimPajama data. This approach offers practical efficiency gains for long-context LLMs and provides a foundation for scaling to larger models and datasets in future work.

Abstract

While Large Language Models (LLMs) have achieved remarkable success in various fields, the efficiency of training and inference remains a major challenge. To address this issue, we propose SUBLLM, short for Subsampling-Upsampling-Bypass Large Language Model, an innovative architecture that extends the core decoder-only framework by incorporating subsampling, upsampling, and bypass modules. The subsampling modules are responsible for shortening the sequence, while the upsampling modules restore the sequence length, and the bypass modules enhance convergence. In comparison to LLaMA, the proposed SUBLLM exhibits significant enhancements in both training and inference speeds as well as memory usage, while maintaining competitive few-shot performance. During training, SUBLLM increases speeds by 26% and cuts memory by 10GB per GPU. In inference, it boosts speeds by up to 37% and reduces memory by 1GB per GPU. The training and inference speeds can be enhanced by 34% and 52% respectively when the context window is expanded to 8192. Our code is available at https://github.com/XiaoMi/subllm.

Figures (5)

• Figure 1: The overall architecture of SUBLLM.
• Figure 2: The impact of various subsampling setups on model performance and speed-up in pre-training. Figure \ref{['fig:sub_r']} illustrates the model with one and two subsampling modules, denoted by (1) and (2), respectively.
• Figure 3: Attention distribution of the 5th block for SUBLLM and the 6th block for LLaMA, where kept indexes in subsampling are highlighted in red.
• Figure 4: A demonstration of the computation of subsampling, upsampling and bypass modules. From bottom to top, the first dashed box represents the learnable subsampling module, the second dashed box represents the upsampling module, and the third dashed box represents the bypass module. Sampling denotes performing sampling with replacement from the weights of the tokens not selected, and these weights are used to be subtracted from the weights of the selected tokens.
• Figure 5: Valid loss comparison of 1.3B models under 4K and 8K Settings.