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

TL;DR

SUBLLM introduces a novel architecture for large language models that significantly improves training and inference efficiency through sequence subsampling, upsampling, and bypass modules, while maintaining competitive performance.

Contribution

The paper presents SUBLLM, an innovative LLM architecture that extends the decoder-only framework with subsampling, upsampling, and bypass modules to enhance efficiency.

Findings

Training speed increased by 26% with 10GB memory reduction per GPU.

Inference speed improved by up to 37% with 1GB memory reduction per GPU.

Efficiency gains are amplified with larger context windows, up to 52% speed increase.

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.