Evaluating Long Range Dependency Handling in Code Generation Models using Multi-Step Key Retrieval

2024-07-23

arXiv.org

AI Summary

This work addresses the bottleneck of large language models (LLMs) in modeling long-range dependencies during code generation, particularly their ability to resolve post-defined function references within realistic 8K-context scenarios. We propose a multi-step key retrieval evaluation framework—more fine-grained than the standard “needle-in-a-haystack” benchmark—and systematically expose a fundamental failure of sliding-window attention mechanisms in cross-window function calls, with failure rates exceeding 90%. Leveraging program call graphs, we further design a lightweight, graph-guided prompting method. Experiments demonstrate that post-definition references degrade model performance by up to 2×; our call-graph-guided prompting improves multi-step retrieval accuracy by up to 3×, substantially mitigating long-range dependency modeling deficiencies.

Technology Category

• Natural Language Processing: Code Generation / Program Synthesis from Natural Language
• Machine Learning: Large Multimodal Models (LMMs)
• Planning, Routing, and Scheduling: Planning with Language Models

Application Category

• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Large language models for search
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs

Abstract

As language models support larger and larger context sizes, evaluating their ability to make effective use of that context becomes increasingly important. We analyze the ability of several code generation models to handle long range dependencies using a suite of multi-step key retrieval tasks in context windows up to 8k tokens in length. The tasks progressively increase in difficulty and allow more nuanced evaluation of model capabilities than tests like the popular needle-in-the-haystack test. We find that performance degrades significantly (up to 2x) when a function references another function that is defined later in the prompt. We also observe that models that use sliding window attention mechanisms have difficulty handling references further than the size of a single window. We perform simple prompt modifications using call graph information to improve multi-step retrieval performance up to 3x. Our analysis highlights different facets of long-context performance and is suggestive of prompt construction strategies for code completion tools

Problem

• Evaluating code generation LLMs' ability to handle long-range dependencies.
• Assessing performance degradation in models when functions reference later-defined functions.
• Improving multi-step retrieval using call graph-based prompt modifications.

Innovation

• Multi-step key retrieval tasks for evaluation
• Sliding window attention mechanism analysis
• Call graph prompt modifications for improvement