H i R o PE : Length Extrapolation for Code Models Using Hierarchical Position (2024)

Abstract: Addressing the limitation of context length in large language models for code-related tasks is the primary focus of this paper. Existing LLMs are constrained by their pre-trained context lengths, leading to performance issues in handling long complex code sequences. Inspired by how human programmers navigate code, we introduce Hierarchical Rotary Position Embedding (HiRoPE), a novel approach that enhances the traditional rotary position embedding into a hierarchical format based on the hierarchical structure of source code. HiRoPE offers easy integration into existing LLMs without extra training costs. Our method is extensively evaluated with various LLMs, demonstrating stable performance in tasks such as language modeling and long code completion. We also introduce a new long code understanding task with real-world code projects, in hopes of promoting further development in this code-related field. Theoretically and experimentally, we find that HiRoPE also addresses the out-of-distribution issue in position encoding. Our HiRoPE significantly expands the context length capabilities of LLMs, enabling inference at lengths exponentially greater than the training length.

## introduction large language models (llms) such as llama-2 @xcite , and codellama @xcite have achieved significant performances in code-related tasks. these transformer-based models excel in code comprehension and generation but face a notable challenge: the limitation of maximum context length. llms are typically pre-trained with a context length ranging from 2k to 16k tokens, which often proves insufficient for complex, extended source code. exceeding this length limitation during inference may lead to performance degradation for these code models, particularly in tasks like project-level code completion or long code generation. ## preliminary we first introduce rotary position embedding in transformer in section 2.1. while existing work usually regards source code as plain text for modeling, we will also introduce the ignored hierarchical information in source code in section 2.2. ## hierarchical rope in this paper, we propose hirope, a hierarchical rotary position encoding for source code modeling. our proposed hirope requires two modified stages: ❶ hierarchical format: we first take the step to transfer the existing rotary position embedding into a hierarchical format. we verify our approach theoretically and find that the hierarchical format can bring stable extrapolation ability for rope. ❷ window mechanism. to ensure performance stability without further training, we also add a window mechanism, so that when dealing with short texts, our proposed method is consistent with the original positional encoding. an illustration of our hirope is shown in figure 2 . ## experiment setup in this section, we aim to answer the following research questions through a series of experiments. details of the evaluation dataset statistics are shown in rq1. how is the language modeling capability of hirope on long code sequences? we evaluate hirope's language modeling ability on codeparrot-valid dataset (codeparrot, 2022) in section 5.1. rq2. how is the language modeling capability of hirope on long natural language sequences? the natural language lacks the explicit hierarchical structure information found in code, so we have made some modifications: we set every 128 tokens as a segment, and encode it as higherlevel position information. we use the evaluation rq3. how does hirope perform in understanding real-world, long-code projects? to evaluate the effect of the method in real long-code scenarios, we design a new evaluation task on real code projects: code symbol understanding in section 5.3. given a long code file, the model is required to output all the function names and class names defined in it. we extract long code files from popular open-sourced code repositories, especially those newly updated code projects to avoid data leakage. details of these code projects are shown in rq4. how does hirope perform on existing benchmarks for long code completion? we further perform the evaluation using two long code completion benchmarks: lcc (guo et al., 2023) and repobench @xcite in section 5.4. ## related work existing large language models are originally trained with fixed context sizes. when dealing with longer sequences, the model's performance may decrease quite drastically. recent studies have explored ways to expand the context length. position interpolation @xcite linearly down-scales the input position indices to match the original context window size of llms with several training steps. similar studies @xcite @xcite also require fine-tuning. however, these methods need additional tuning in longer contexts or face a disastrous collapse after the extrapolation bound. there are also some approaches without training. some studies use window attention to clip the long sequences such as @xcite @xcite ). however, these methods rely on local information and may not effectively expand the context window, struggling with long dependencies in code. recently, some methods have explored modifying the relative distance to extend the extrapolation length @xcite @xcite and focus on the natural language text. we pursue the research line of training-free methods and propose considering the structural information of the code when modeling the position. our work is inspired by the hierarchical embedding in our former work @xcite . we notice hierarchical positions are used in some concurrent work @xcite in nlp and cv domains. we expand the traditional rope method into a hierarchical format and prove the effectiveness of our hirope through theoretical derivations and practical experiments. ## conclusion we propose hirope, a training-free solution to the context length limitation in llms for long code modeling. we integrate the hierarchical structure of source code into position encoding of llms. experiments demonstrate that hirope achieves stable improvements on diverse code-related tasks. our work not only addresses a critical limitation in llm applications but also opens new avenues for long structured data modeling research.