arXiv Machine Learning
Sep 17

Long-Context Demonstration Selection Using State Space Models

The paper addresses the challenge of selecting demonstrations for long-context language model queries, where transformer inference costs grow quadratically with sequence length. It proposes two algorithms that distill transformer behavior into state space models (SSMs) with linear inference time, partitioning transformer layers into groups and estimating separate SSMs for each. The distilled SSMs achieve less than 0.7% approximation error, and in downstream tasks they reduce FLOPs by 14.2× while improving accuracy by 6.48% compared to baseline methods.

By Ziniu Zhang, Zhenshuo Zhang, Ruoxuan Xiong, Gene Cooperman, Hongyang R. Zhang