Efficient Beam Search for Large Language Models Using Trie-Based Decoding

• URL: http://arxiv.org/abs/2502.00085v1
• Date: Fri, 31 Jan 2025 16:22:36 GMT
• Title: Efficient Beam Search for Large Language Models Using Trie-Based Decoding
• Authors: Brian J Chan, Jui-Hung Cheng, Mao Xun Huang, Chao-Ting Chen, Hen-Hsen Huang,
• Abstract summary: We introduce a novel parallel decoding method that addresses the memory inefficiency of batch-based beam search.<n>By sharing a single cache among all beams that share the same prefix, the proposed method not only reduces memory consumption dramatically but also enables parallel decoding across all branches.<n>This innovative use of a prefix tree offers an efficient alternative for beam search, achieving significant memory savings while preserving inference speed, making it particularly well-suited for memory-constrained environments or large-scale model deployments.
• Score: 10.302821791274129
• License: http://creativecommons.org/licenses/by-sa/4.0/
• Abstract: In Transformer-based sequence-to-sequence generation, beam search has proven effective in enhancing the quality of generated sequences compared to greedy decoding. Conventional beam search methods typically adopt either a sequential or batch-based approach. The sequential approach, while memory-efficient, requires multiple decoding passes to construct a complete search tree, leading to significantly slower inference. On the other hand, the batch-based approach enables parallel computation across beams, but at the expense of high memory consumption due to the need to maintain separate key-value (KV) caches for each beam. In this study, we introduce a novel trie (prefix-tree)-based parallel decoding method that addresses the memory inefficiency of batch-based beam search. By sharing a single KV cache among all beams that share the same prefix, the proposed method not only reduces memory consumption dramatically but also enables parallel decoding across all branches. This innovative use of a prefix tree offers an efficient alternative for beam search, achieving significant memory savings while preserving inference speed, making it particularly well-suited for memory-constrained environments or large-scale model deployments.

Related papers

• From Verbatim to Gist: Distilling Pyramidal Multimodal Memory via Semantic Information Bottleneck for Long-Horizon Video Agents [78.30630000529133]
  We propose MM-Mem, a pyramidal multimodal memory architecture grounded in Fuzzy-Trace Theory.<n> MM-Mem memory structures hierarchically into a Sensory Buffer, Episodic Stream, and Symbolic.<n>Experiments confirm the effectiveness of MM-Mem on both offline and streaming tasks.
  arXiv  Detail & Related papers  (2026-03-02T05:12:45Z)
• KEEP: A KV-Cache-Centric Memory Management System for Efficient Embodied Planning [8.216400469571084]
  We propose KEEP, a KV-cache-centric memory management system for efficient embodied planning.<n>KEEP features 3 key innovations: (1) a Static-Dynamic Memory Construction algorithm that reduces KV cache recomputation by mixed-granularity memory group; (2) a Multi-hop Memory Re-computation algorithm that dynamically identifies important cross-attention among different memory groups; and (3) a Layer-balanced Memory Loading that eliminates unbalanced KV cache loading and cross-attention across different layers.
  arXiv  Detail & Related papers  (2026-02-27T01:48:07Z)
• SimpleMem: Efficient Lifelong Memory for LLM Agents [73.74399447715052]
  We introduce SimpleMem, an efficient memory framework based on semantic lossless compression.<n>We propose a three-stage pipeline designed to maximize information density and token utilization.<n> Experiments on benchmark datasets show that our method consistently outperforms baseline approaches in accuracy, retrieval efficiency, and inference cost.
  arXiv  Detail & Related papers  (2026-01-05T21:02:49Z)
• Remember Me, Refine Me: A Dynamic Procedural Memory Framework for Experience-Driven Agent Evolution [52.76038908826961]
  We propose $textbfReMe$ ($textitRemember Me, Refine Me$) to bridge the gap between static storage and dynamic reasoning.<n>ReMe innovates across the memory lifecycle via three mechanisms: $textitmulti-faceted distillation$, which extracts fine-grained experiences.<n>Experiments on BFCL-V3 and AppWorld demonstrate that ReMe establishes a new state-of-the-art in agent memory system.
  arXiv  Detail & Related papers  (2025-12-11T14:40:01Z)
• Leveraging Lightweight Generators for Memory Efficient Continual Learning [0.01874930567916036]
  Catastrophic forgetting can be trivially alleviated by keeping all data from previous tasks in memory.<n>This paper aims to decrease required memory for memory-based continuous learning algorithms.
  arXiv  Detail & Related papers  (2025-06-24T14:59:52Z)
• Memory-Efficient FastText: A Comprehensive Approach Using Double-Array Trie Structures and Mark-Compact Memory Management [0.0]
  FastText has established itself as a fundamental algorithm for learning word representations.<n>However, its hash-based bucketing mechanism introduces critical limitations for large-scale industrial deployment.<n>This paper presents a memory optimization framework that reimagines FastText's memory management.
  arXiv  Detail & Related papers  (2025-06-02T02:11:22Z)
• A Universal Framework for Compressing Embeddings in CTR Prediction [68.27582084015044]
  We introduce a Model-agnostic Embedding Compression (MEC) framework that compresses embedding tables by quantizing pre-trained embeddings.<n>Our approach consists of two stages: first, we apply popularity-weighted regularization to balance code distribution between high- and low-frequency features.<n> Experiments on three datasets reveal that our method reduces memory usage by over 50x while maintaining or improving recommendation performance.
  arXiv  Detail & Related papers  (2025-02-21T10:12:34Z)
• PrefixKV: Adaptive Prefix KV Cache is What Vision Instruction-Following Models Need for Efficient Generation [65.36715026409873]
  Key-value (KV) cache, necessitated by the lengthy input and output sequences, notably contributes to the high inference cost.<n>We present PrefixKV, which reframes the challenge of determining KV cache sizes for all layers into the task of searching for the optimal global prefix configuration.<n>Our method achieves the state-of-the-art performance compared with others.
  arXiv  Detail & Related papers  (2024-12-04T15:48:59Z)
• Dynamic-Width Speculative Beam Decoding for Efficient LLM Inference [35.730941605490194]
  Large language models (LLMs) have shown outstanding performance across numerous real-world tasks. Speculative decoding has emerged as a promising solution, leveraging a smaller auxiliary model to draft future tokens. This paper explores the novel integration of speculative decoding with beam sampling.
  arXiv  Detail & Related papers  (2024-09-25T02:20:42Z)
• An Efficient Procedure for Computing Bayesian Network Structure Learning [0.9208007322096532]
  We propose a globally optimal Bayesian network structure discovery algorithm based on a progressively leveled scoring approach. Experimental results indicate that our method, when using only memory, not only reduces peak memory usage but also improves computational efficiency.
  arXiv  Detail & Related papers  (2024-07-24T07:59:18Z)
• Simple linear attention language models balance the recall-throughput tradeoff [60.06020449520365]
  We propose BASED, a simple architecture combining linear and sliding window attention.<n>We train language models up to 1.3b parameters and show that BASED matches the strongest sub-quadratic models in perplexity and outperforms them on real-world recall-intensive tasks by 6.22 accuracy points.
  arXiv  Detail & Related papers  (2024-02-28T19:28:27Z)
• MAMBA: Multi-level Aggregation via Memory Bank for Video Object Detection [35.16197118579414]
  We propose a multi-level aggregation architecture via memory bank called MAMBA. Specifically, our memory bank employs two novel operations to eliminate the disadvantages of existing methods. Compared with existing state-of-the-art methods, our method achieves superior performance in terms of both speed and accuracy.
  arXiv  Detail & Related papers  (2024-01-18T12:13:06Z)
• HashReID: Dynamic Network with Binary Codes for Efficient Person Re-identification [3.3372444460738357]
  Biometric applications, such as person re-identification (ReID), are often deployed on energy constrained devices. While recent ReID methods prioritize high retrieval performance, they often come with large computational costs and high search time. We propose an input-adaptive network with multiple exit blocks, that can terminate early if the retrieval is straightforward or noisy.
  arXiv  Detail & Related papers  (2023-08-23T04:01:54Z)
• Efficient Nearest Neighbor Search for Cross-Encoder Models using Matrix Factorization [60.91600465922932]
  We present an approach that avoids the use of a dual-encoder for retrieval, relying solely on the cross-encoder. Our approach provides test-time recall-vs-computational cost trade-offs superior to the current widely-used methods.
  arXiv  Detail & Related papers  (2022-10-23T00:32:04Z)
• Asymmetric Scalable Cross-modal Hashing [51.309905690367835]
  Cross-modal hashing is a successful method to solve large-scale multimedia retrieval issue. We propose a novel Asymmetric Scalable Cross-Modal Hashing (ASCMH) to address these issues. Our ASCMH outperforms the state-of-the-art cross-modal hashing methods in terms of accuracy and efficiency.
  arXiv  Detail & Related papers  (2022-07-26T04:38:47Z)
• Recall@k Surrogate Loss with Large Batches and Similarity Mixup [62.67458021725227]
  Direct optimization, by gradient descent, of an evaluation metric is not possible when it is non-differentiable. In this work, a differentiable surrogate loss for the recall is proposed. The proposed method achieves state-of-the-art results in several image retrieval benchmarks.
  arXiv  Detail & Related papers  (2021-08-25T11:09:11Z)
• Semantically Constrained Memory Allocation (SCMA) for Embedding in Efficient Recommendation Systems [27.419109620575313]
  A key challenge for deep learning models is to work with millions of categorical classes or tokens. We propose a novel formulation of memory shared embedding, where memory is shared in proportion to the overlap in semantic information. We demonstrate a significant reduction in the memory footprint while maintaining performance.
  arXiv  Detail & Related papers  (2021-02-24T19:55:49Z)
• Decoupled and Memory-Reinforced Networks: Towards Effective Feature Learning for One-Step Person Search [65.51181219410763]
  One-step methods have been developed to handle pedestrian detection and identification sub-tasks using a single network. There are two major challenges in the current one-step approaches. We propose a decoupled and memory-reinforced network (DMRNet) to overcome these problems.
  arXiv  Detail & Related papers  (2021-02-22T06:19:45Z)
• A Genetic Algorithm for Obtaining Memory Constrained Near-Perfect Hashing [0.0]
  We present an approach based on hash tables that focuses on both minimizing the number of comparisons performed during the search and minimizing the total collection size. The paper results show that near-perfect hashing is faster than binary search, yet uses less memory than perfect hashing.
  arXiv  Detail & Related papers  (2020-07-16T12:57:15Z)
• Best-First Beam Search [78.71330480725668]
  We show that the standard implementation of beam search can be made up to 10x faster in practice. We propose a memory-reduced variant of Best-First Beam Search, which has a similar beneficial search bias in terms of downstream performance.
  arXiv  Detail & Related papers  (2020-07-08T05:56:01Z)
• Reinforcing Short-Length Hashing [61.75883795807109]
  Existing methods have poor performance in retrieval using an extremely short-length hash code. In this study, we propose a novel reinforcing short-length hashing (RSLH) In this proposed RSLH, mutual reconstruction between the hash representation and semantic labels is performed to preserve the semantic information. Experiments on three large-scale image benchmarks demonstrate the superior performance of RSLH under various short-length hashing scenarios.
  arXiv  Detail & Related papers  (2020-04-24T02:23:52Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.