C2LLM Technical Report: A New Frontier in Code Retrieval via Adaptive Cross-Attention Pooling

🍞 Hook: Imagine you’re searching a giant library for one perfect recipe. You could skim every page equally, or you could hunt for the parts with ingredient lists and cooking steps. Which is better?

🥬 The Concept (Code Retrieval): It’s the job of finding the most relevant piece of code for a natural-language request. How it works:

1. Turn the query (text) and code (tokens) into vectors called embeddings.
2. Store code embeddings in a database.
3. Turn the user’s question into an embedding and find nearest neighbors. Why it matters: Without good embeddings, the search feels random—relevant code hides among millions of files. 🍞 Anchor: Ask, “open a jsonl file in Python and read all lines,” and the system should pull up a snippet using json, open, and iteration.

🍞 Hook: You know how a nickname captures the essence of a person in one word?

🥬 The Concept (Embedding): An embedding is a single vector that summarizes a whole sequence (like code or a sentence). How it works:

1. An LLM reads tokens and builds internal features.
2. A pooling method compresses them into one vector.
3. That vector is used for fast similarity search. Why it matters: If the summary misses key details, the search engine won’t find the right code. 🍞 Anchor: A function that sorts numbers should have an embedding close to “sort a list in ascending order.”

🍞 Hook: Picture averaging all your school grades—even the practice quizzes—into one number. It’s fair, but it hides which tests mattered most.

🥬 The Concept (Mean Pooling): Mean pooling averages information from all token positions to form one vector. How it works:

1. Collect token features from the model’s last layer.
2. Add them up and divide by the count.
3. Use that average as the sequence embedding. Why it matters: For code, not all tokens are equally useful; averaging can wash out key lines like function signatures. 🍞 Anchor: In a 300-line file, averaging treats comments and imports the same as the core algorithm—making the final vector blurry.

🍞 Hook: Think of stuffing all your notes into the last page of your notebook—everything’s squeezed there.

🥬 The Concept (EOS Token Representation): Using the final token’s vector as the whole sequence’s embedding. How it works:

1. Process tokens left-to-right.
2. Take the last token (EOS) feature.
3. Use it as the sequence vector. Why it matters: This creates a bottleneck—long code with rich structure gets crammed into one spot. 🍞 Anchor: A big Python file’s final token can’t carry all the logic about classes, loops, and I/O.

🍞 Hook: Imagine reading a story where each line depends on the previous one—like a chain of clues.

🥬 The Concept (Causal Representations): They’re features learned by LLMs that read left-to-right, where each token can only see the past. How it works:

1. The model scans tokens in order.
2. Each new token summarizes what came before.
3. The chain builds meaningful, time-aware features. Why it matters: Many top code LLMs use causal training; breaking this setup can waste their strengths. 🍞 Anchor: A function’s return line makes sense only after reading its variable definitions above.

The world before: Developers relied on mean pooling (blurry summaries) or EOS (squeezed summaries). These methods were designed around text settings, not massive, structured code. Code files are long and packed with meaning in specific places (function names, type hints, loops).

The problem: Embedding all that code into one vector without losing the most important bits. Mean pooling over-smooths; EOS bottlenecks; both miss the structure that matters to retrieval and agents.

Failed attempts:

• Mean pooling with bidirectional encoders clashed with causal code LLMs and underused their pretraining strengths.
• EOS compressed everything into one spot.
• Latent attention methods still needed mean pooling afterward, keeping the same weakness.

The gap: A pooling method that (1) respects causal LLMs, (2) looks at every token and picks the important ones, and (3) flexibly shrinks embedding size for vector databases without extra complex training objectives.

Real stakes: Better code retrieval speeds up engineers, powers code agents that plan-search-edit, and helps answer tricky, multi-turn questions mixing text and code. In daily life, this means fewer hours hunting StackOverflow, faster bug fixes, and safer, more consistent code reuse across big teams.

🍞 Hook: You know how a great teacher doesn’t average every sentence you say—they hone in on the key points you made and summarize those?

🥬 The Concept (Pooling by Multihead Attention, PMA): PMA is a tiny cross-attention layer that uses one learnable query to attend to all token features and produce one smart embedding. How it works:

1. Feed code into a causal LLM to get token features.
2. Use a learned query vector to “look at” all tokens via cross-attention.
3. Mix the important bits into a single vector; normalize and output. Why it matters: It keeps the LLM’s strengths, avoids averaging away key info, removes the EOS bottleneck, and can change embedding size flexibly. 🍞 Anchor: Given a long function, PMA will focus more on the signature and core loop than on comments or imports.

Aha! Moment in one sentence: Instead of averaging or squeezing, use a single, learned spotlight (cross-attention) to gather exactly the right pieces from every token into one powerful vector.

Three analogies:

• Microscope analogy: PMA is a microscope you point at the parts of code that matter most, then you snap one clear photo (the embedding).
• Chef analogy: PMA tastes every ingredient (token) but only keeps the strongest flavors for the final dish (the vector).
• Librarian analogy: PMA scans the whole book, bookmarks the crucial pages, and writes a sharp summary.

Before vs After:

• Before: Mean pooling washes out signals; EOS bottlenecks; dimension tied to the LLM’s hidden size or costly tricks like MRL needed.
• After: Cross-attention pooling targets key tokens, scales across code lengths, and decouples LLM hidden size from final embedding size—handy for vector DBs.

🍞 Hook: Imagine asking a class of students for ideas, but instead of averaging every opinion, you ask one focused question and listen most closely to the most relevant answers.

🥬 The Concept (Cross-Attention): Cross-attention lets one sequence (a query) decide which parts of another sequence (keys/values) to focus on. How it works:

1. Create a query vector.
2. Compare it to keys made from token features.
3. Use the scores to weight values (the same token features) and blend them. Why it matters: It gives the model control to emphasize function names, signatures, or critical logic over noise. 🍞 Anchor: When searching for “SQL injection sanitize,” cross-attention highlights sanitizer calls and parameterized queries.

🍞 Hook: Think of swapping your heavy backpack for a neat pencil case that still holds everything important.

🥬 The Concept (Embedding Dimension Adaptation): It’s the ability to output a smaller embedding size than the LLM’s hidden dimension without special training. How it works:

1. Project token features into a target dimension.
2. Run cross-attention there.
3. Output the compact vector. Why it matters: Smaller vectors are cheaper to store and search in big databases, with minimal performance loss. 🍞 Anchor: From a 4096-d hidden state to a 1024-d embedding that’s fast to index and retrieve.

Why it works (intuition): Code meaning is concentrated in specific regions (e.g., definitions, control flow), not uniformly spread. A learned query acts like a smart magnet that pulls the most relevant iron filings from the whole file. Because the LLM’s causal training already builds rich, step-by-step features, PMA simply harvests them efficiently, preserving structure while preventing dilution or bottlenecks.

Building blocks:

• Causal LLM backbone supplies high-quality token features.
• PMA (single-learned-query cross-attention) selects and mixes salient tokens.
• Lightweight projections decouple hidden and embedding dimensions.
• Simple MLP + LayerNorm stabilize and refine the pooled vector.
• Contrastive training shapes embeddings so matching query–code pairs become near neighbors.

🍞 Hook: Imagine a smart vacuum that only picks up the dirt, not your toys.

🥬 The Concept (Adaptive Cross-Attention Pooling): It’s pooling that learns where to focus across all tokens, adaptively weighting the most important ones. How it works:

1. Feed tokens to an LLM to get features.
2. Use a learnable query to score tokens.
3. Blend high-scoring tokens into one vector. Why it matters: Without adaptive pooling, long code gets misrepresented or flattened. 🍞 Anchor: In a codebase, it helps a “file upload” query find functions that validate file types and size limits, not just any mention of “file.”

At a high level: Input (query or code) → LLM encodes tokens → PMA cross-attention pools to one vector (and adapts dimension) → Normalize/MLP → Embedding → Contrastive training to align queries and code.

Step 1: Tokenize and encode with a causal code LLM

• What happens: The input (text query or code) is turned into tokens, then passed through a pretrained Qwen2.5-Coder (0.5B or 7B) to produce last-layer hidden states for each token.
• Why this step exists: The LLM builds rich, position-aware features that capture syntax and semantics—raw tokens aren’t meaningful enough.
• Example: For “open a jsonl file and read lines,” tokens around open, with, json, and iteration patterns get distinctive features.

🍞 Hook: Think of reading a book page by page; each new page makes sense because you’ve read the earlier ones.

🥬 The Concept (Causal Representations): They’re step-by-step features built by reading left-to-right, where each position summarizes what came before. How it works:

1. Mask future tokens.
2. Each token sees only the past.
3. Features grow cumulatively. Why it matters: Preserves the strengths of code LLM pretraining and keeps temporal order. 🍞 Anchor: A variable used in line 120 is meaningful because it was defined in line 30.

Step 2: Pool with PMA (Pooling by Multihead Attention)

• What happens: A single learnable query looks at all token features via cross-attention and blends the most relevant information into one vector. Multi-heads let it capture different signals (e.g., signature, loop, API usage) in parallel.
• Why this step exists: It avoids the two classic failures: over-smoothing (mean) and bottlenecking (EOS). It also lets us pick a smaller output dimension.
• Example: In a function for CSV parsing, PMA may emphasize delimiter handling, error checks, and return type.

🍞 Hook: You know how when you ask a precise question, you notice the most relevant parts of an answer first?

🥬 The Concept (Cross-Attention): It’s a way to compare a query to many tokens and pick the best matches. How it works:

1. Project the query and tokens.
2. Score how well each token matches the query.
3. Blend tokens by those scores to get one summary. Why it matters: It puts attention where it counts, not evenly everywhere or only at the end. 🍞 Anchor: For “sanitize user input,” tokens with regex checks and escaping functions score higher.

Step 3: Dimension adaptation and refinement

• What happens: Projections shrink features to the target embedding size; the output passes through LayerNorm and a small MLP for stability and expressiveness.
• Why this step exists: Embeddings need to be compact for fast vector search; normalization stabilizes training and retrieval quality.
• Example: Go from 3584-d hidden states to a 1024-d embedding suitable for Milvus/FAISS.

🍞 Hook: Imagine fitting your suitcase to airline rules without leaving behind the essentials.

🥬 The Concept (Embedding Dimension Adaptation): Outputting a smaller, efficient vector than the LLM’s hidden size. How it works:

1. Linear layers reduce dimensions.
2. PMA runs in the reduced space.
3. LayerNorm/MLP polish the result. Why it matters: Saves storage and speeds retrieval in large-scale systems. 🍞 Anchor: A company with billions of code snippets can cut costs by using 768–1024-d embeddings instead of very large ones.

Step 4: Train with contrastive learning

• What happens: The model learns to bring matched pairs (query, correct code) closer and push mismatched ones apart. It uses in-batch negatives and 7 hard negatives per query, with temperature τ=0.05.
• Why this step exists: Retrieval is a ranking problem; contrastive loss directly shapes the space for nearest-neighbor search.
• Example: “merge two sorted arrays” pulls close to code with two-pointer merge; pushes away unrelated UI code.

🍞 Hook: Spot-the-difference puzzles train your eyes to see what matches and what doesn’t.

🥬 The Concept (Contrastive Learning): A training method that compares examples to learn similarities and differences. How it works:

1. Encode queries and codes.
2. Pull true pairs together.
3. Push false pairs apart (especially hard negatives). Why it matters: It sculpts the embedding space so the right answers pop to the top. 🍞 Anchor: Photos of the same dog breed cluster together; different breeds spread apart—same idea for code tasks.

🍞 Hook: Practicing basketball against tougher defenders makes you sharper.

🥬 The Concept (Hard Negatives): These are tricky wrong answers that look right, forcing the model to learn fine distinctions. How it works:

1. For each query, pick 7 similar-but-wrong snippets.
2. Penalize the model if it gets fooled.
3. Repeat across batches to expand negative variety. Why it matters: Prevents the model from taking shortcuts and strengthens real-world retrieval. 🍞 Anchor: “parse JSON lines” vs “parse CSV lines” are hard negatives that teach careful attention to file format.

Training efficiency details

• LoRA (rank 64, alpha 32) fine-tunes large backbones with small parameter updates, saving memory and time.
• FlashAttention 2 accelerates attention computations, enabling longer sequences and larger batches.
• Left-padding to length 1024 standardizes training.
• Weighted checkpoint merging stabilizes the final model.

🍞 Hook: Instead of rebuilding a car’s engine, you just swap in a small turbo that boosts speed.

🥬 The Concept (LoRA): A fine-tuning trick that adds small low-rank adapters to a big model so you train far fewer weights. How it works:

1. Freeze most of the LLM.
2. Insert tiny adapter layers.
3. Train just those adapters. Why it matters: You get strong performance without huge compute costs. 🍞 Anchor: Fine-tuning a 7B model on a single high-end GPU becomes feasible.

🍞 Hook: Imagine a super-efficient highway that keeps traffic moving even at rush hour.

🥬 The Concept (FlashAttention 2): A fast attention algorithm that reduces memory use and speeds up training/inference. How it works:

1. Reorder attention math to minimize memory reads/writes.
2. Parallelize better across GPUs.
3. Keep numerical stability. Why it matters: Lets you train with longer sequences and bigger batches affordably. 🍞 Anchor: Encoding 1k-token files smoothly instead of constantly hitting memory limits.

Secret sauce

• The single-learned-query PMA turns the LLM’s causal features into a precise, compact summary.
• This pooling simultaneously solves information selection and dimension control.
• Together with contrastive training and hard negatives, it creates embeddings that excel at real code retrieval.

The test: The team evaluated on 12 retrieval tasks in MTEB-Code, which mix classic code search, code-to-text, text-to-code, multi-turn Q&A, and cross-language translation of code. This suite measures not just basic matching but also complex reasoning and understanding of intent.

🍞 Hook: Think of a big school tournament with many events—math, spelling, debate—so you can’t win by being good at only one thing.

🥬 The Concept (MTEB-Code Benchmark): A collection of diverse retrieval tasks used to score embedding models on code search and related challenges. How it works:

1. Prebuild a database of code/text items for each task.
2. Encode queries and candidates into embeddings.
3. Rank by similarity; compute metrics; average across tasks. Why it matters: It’s a trusted scoreboard that shows who’s best overall and who handles tricky, real tasks. 🍞 Anchor: If your model shines on CodeFeedback and CodeSearchNet, it likely serves developers well in practice.

The competition: Strong baselines include Qwen3-Embed (4B/8B), Seed1.6-Embed (closed-source), INF-Retriever (1.5B/7B), and EmbeddingGemma (0.3B). Many are general-purpose text embedders; a few are code-focused but not all appear on MTEB-Code.

Scoreboard highlights (with context):

• C2LLM-7B: 80.75 average—1st overall. Think of it as getting an A+ when top classmates are at A or A−.
• Seed1.6-Embed: 80.71—neck-and-neck but edged out.
• Qwen3-Embed-8B: 80.69—slightly below C2LLM-7B despite having a similar or larger backbone.
• C2LLM-0.5B: 75.46—best under 1B parameters; beats INF-Retriever-7B (69.70), like a lightweight athlete outpacing a heavyweight.

Task-level notes:

• CodeFeedback (multi-turn 94.32, single-turn 90.66): Particularly strong—suggests PMA captures intent across mixed text+code conversations.
• CodeSearchNet (e.g., CSN/CCR/CoIR near or above 90s): Solid classical code search.
• Cross-language tasks (CodeTransOcean): Competitive, indicating the pooling stays robust when semantics shift across languages.

Surprising findings:

• The single-query PMA, though tiny compared to the LLM, delivers outsized gains—showing that where and how you pool can matter as much as the backbone size.
• The 0.5B model’s efficiency: With far fewer parameters, it stays strong on diverse tasks, implying cross-attention pooling scales down gracefully.
• Complex reasoning tasks benefit most—exactly where mean/EOS pooling struggles due to dispersed signals across long inputs.

Takeaway: On a challenging, public leaderboard, C2LLM’s PMA-based pooling proves its value across sizes, not just at the biggest scale.

Limitations:

• Single-query PMA may underrepresent documents with multiple, equally important themes; a multi-query variant might help.
• Max sequence length (1024 during training) can clip very long files; extreme-long contexts may need chunking or specialized encoders.
• Performance depends on training data coverage; rare languages, frameworks, or enterprise-specific APIs may see weaker results.
• PMA adds a (small) extra layer; while negligible vs. 7B parameters, it’s still another component to maintain.
• Contrastive training quality hinges on good hard negatives; poor mining can cap gains.

Required resources:

• For 7B: At least a modern GPU with enough memory for inference; multi-GPU recommended for training even with LoRA.
• Vector database capable of handling millions to billions of embeddings efficiently (e.g., FAISS/Milvus).
• Data prep pipelines for tokenization, deduplication, and chunking long files.

When NOT to use:

• Ultra-low-latency, ultra-tiny edge devices where any LLM backbone is too heavy.
• Tasks requiring full-program static analysis or graph reasoning (e.g., inter-file dependency resolution) beyond what embeddings can capture.
• Domains with zero overlap to training distribution and no fine-tuning budget (consider domain adaptation first).

Open questions:

• Would multiple learnable queries (multi-vector pooling) further boost multi-topic files?
• How far can dimension reduction go before retrieval drops—what’s the optimal sweet spot per domain?
• Can PMA help repository-level retrieval (multi-file) or agent planning memory?
• How to best mine and schedule hard negatives automatically at scale?
• Extending to massively multilingual codebases and specialized DSLs while preserving compact vectors.

Three-sentence summary: C2LLM introduces a tiny but powerful PMA pooling layer on top of causal code LLMs to build better code embeddings. By learning to attend to the most important tokens across an entire file—and by flexibly shrinking embedding size—C2LLM avoids the pitfalls of mean and EOS pooling. The result is state-of-the-art code retrieval on MTEB-Code, with strong performance even at the 0.5B scale.

Main achievement: Showing that a single-learned-query cross-attention pooling can unlock the potential of causal code LLMs for retrieval, beating larger or closed models on a public leaderboard.

Future directions: Explore multi-query pooling for multi-topic files, push dimension reduction farther without accuracy loss, scale to repository-level retrieval, and extend to more languages and domains.

Why remember this: In embeddings, the pooling choice can matter as much as the model itself—PMA is a simple, general tool that preserves crucial information, scales across sizes, and makes real developer workflows faster and smarter.