MemoSight: Unifying Context Compression and Multi Token Prediction for Reasoning Acceleration

TL;DR

MemoSight is a unified framework that enhances reasoning efficiency in large language models by combining context compression and multi-token prediction, significantly reducing memory usage and speeding up inference.

Contribution

It introduces a novel, minimalist design for integrating context compression and multi-token prediction, improving reasoning speed and memory efficiency without sacrificing performance.

Findings

Reduces KV cache footprint by up to 66%

Speeds up inference by 1.56x

Outperforms existing CoT compression methods

Abstract

While Chain-of-thought (CoT) reasoning enables LLMs to solve challenging reasoning problems, as KV cache grows linearly with the number of generated tokens, CoT reasoning faces scaling issues in terms of speed and memory usage. In this work, we propose MemoSight (Memory-Foresight-based reasoning), a unified framework that integrates both context compression and multi-token prediction to mitigate the efficiency issues while maintaining CoT reasoning performance. Our framework adopts the same minimalist design for both context compression and multi-token prediction via special tokens and their corresponding position layout tailored to each token type. Comprehensive experiments on four reasoning benchmarks demonstrate that MemoSight reduces the KV cache footprint by up to 66% and accelerates inference by 1.56x, while outperforming existing CoT compression methods.