Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

LLavaCode: Compressed Code Representations for Retrieval-Augmented Code Generation

Daria Cherniuk, Nikita Sukhorukov, Nikita Sushko, Daniil Gusak, Danil Sivtsov, Elena Tutubalina, Evgeny Frolov

TL;DR

The paper tackles latency bottlenecks in retrieval-augmented code generation for interactive IDEs by compressing retrieved code into compact, semantically meaningful embeddings using a LLaVA-like projection. It introduces LlavaCode, a single-stage framework where only a lightweight projector is trained to map encoder embeddings into the code-LLM embedding space, with a composite loss that combines cross-entropy, REINFORCE-based sequence rewards (EM and ES), and a cosine-alignment term to prevent representation collapse, while keeping the encoder and reader LLM frozen. Empirical results on the Python subset of The Stack show 20-38% reductions in TTFT and improvements in EM/ES over baselines with negligible latency increases, and ablations demonstrate the benefit of the combined loss and the detrimental effect of KL-loss in this setting. The work also explores code modalities like ASTs and encoder choices, and discusses deployment scenarios (disaggregated vs colocated) and future directions, including RL variants and language expansion.

Abstract

Retrieval-augmented generation has emerged as one of the most effective approaches for code completion, particularly when context from a surrounding repository is essential. However, incorporating context significantly extends sequence length, leading to slower inference - a critical limitation for interactive settings such as IDEs. In this work, we introduce LlavaCode, a framework that compresses code into compact, semantically rich representations interpretable by code LLM, enhancing generation quality while reducing the retrieved context to only a few compressed single-token vectors. Using a small projector module we can significantly increase the EM and ES metrics of coding model with negligible latency increase. Our experiments demonstrate that compressed context enables 20-38% reduction in Time-to-First-Token (TTFT) on line completion tasks compared to full-RAG pipelines.

LLavaCode: Compressed Code Representations for Retrieval-Augmented Code Generation

TL;DR

The paper tackles latency bottlenecks in retrieval-augmented code generation for interactive IDEs by compressing retrieved code into compact, semantically meaningful embeddings using a LLaVA-like projection. It introduces LlavaCode, a single-stage framework where only a lightweight projector is trained to map encoder embeddings into the code-LLM embedding space, with a composite loss that combines cross-entropy, REINFORCE-based sequence rewards (EM and ES), and a cosine-alignment term to prevent representation collapse, while keeping the encoder and reader LLM frozen. Empirical results on the Python subset of The Stack show 20-38% reductions in TTFT and improvements in EM/ES over baselines with negligible latency increases, and ablations demonstrate the benefit of the combined loss and the detrimental effect of KL-loss in this setting. The work also explores code modalities like ASTs and encoder choices, and discusses deployment scenarios (disaggregated vs colocated) and future directions, including RL variants and language expansion.

Abstract

Retrieval-augmented generation has emerged as one of the most effective approaches for code completion, particularly when context from a surrounding repository is essential. However, incorporating context significantly extends sequence length, leading to slower inference - a critical limitation for interactive settings such as IDEs. In this work, we introduce LlavaCode, a framework that compresses code into compact, semantically rich representations interpretable by code LLM, enhancing generation quality while reducing the retrieved context to only a few compressed single-token vectors. Using a small projector module we can significantly increase the EM and ES metrics of coding model with negligible latency increase. Our experiments demonstrate that compressed context enables 20-38% reduction in Time-to-First-Token (TTFT) on line completion tasks compared to full-RAG pipelines.
Paper Structure (23 sections, 5 equations, 4 figures, 9 tables)

This paper contains 23 sections, 5 equations, 4 figures, 9 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Coding LLMs
  4. Context compression methods
  5. Embedding models for code
  6. Reinforcement Learning in Language Modeling
  7. Methodology
  8. Model architecture
  9. Cross-Entropy Issue
  10. REINFORCE
  11. Cosine Alignment Loss
  12. Final Loss Function
  13. KL-loss
  14. Experiments
  15. Dataset
  16. ...and 8 more sections

Figures (4)

  • Figure 1: Comparison between Vanilla RAG \ref{['fig:architecture_rag']} and LlavaCode \ref{['fig:architecture_llavacode']} architectures. Instead of retrieving text passages and putting them into the context of the reader language model, LlavaCode uses a pretrained encoder to compress the text representations and projects them into continuous tokens, thus, reducing the prompt processing time.
  • Figure 2: Pairwise cosine distances between vector outputs. While the encoder representations remain well-separated (a), the projected vectors may collapse, becoming nearly indistinguishable (b). Introducing the Cosine Alignment Loss \ref{['eq:cos_loss']} helps preserve the distinctions among the projections, preventing excessive overlap.
  • Figure 3: Relationship between KL-divergence loss and performance metrics (Exact Match (EM) and Edit Similarity (ES)) on Qwen2.5-Coder-1.5B. The figure shows that reduction in KL-loss is accompanied by decreases in averaged Exact Match (EM) and Edit Similarity (ES) metrics.
  • Figure 4: Relationship between the three loss components (Cross-Entropy, REINFORCE, and Cosine Alignment) and the evaluation metrics Exact Match (EM) and Edit Similarity (ES).