Focus-Guided Transformers: Selective Context Attention for Efficient Code-Agent Reasoning

Summary. FGT computes a three-component relevance signal - hidden-state probing, TF-IDF keyword matching, and structural code-pattern detection - to select the most informative subset of context chunks before generation. Across two codebases totalling 136 files and 261K tokens, evaluated at four context sizes (4.4K to 9.1K tokens), it achieves 70–91% context reduction while maintaining or improving answer quality on 67% of queries.

We identify a critical threshold-minimum tradeoff: aggressive filtering (threshold > 0.05) risks discarding cross-file dependencies, while permissive filtering with a chunk floor (threshold 0.02, min 5 chunks) improves cross-file query accuracy by up to 481 points over baseline. We also show that unguided baseline quality degrades 7–10% per doubling of context length, while FGT quality remains stable - suggesting that selective context becomes increasingly important at the 128K–1M token scales typical of production coding agents.

1. Introduction

Coding agents built on large language models (LLMs) are transforming software development workflows. These agents ingest entire codebases as context, reasoning across files, functions, and architectural patterns to answer questions, generate code, and debug issues. The dominant paradigm feeds all retrieved context into the model's attention mechanism, treating every token as equally informative.

This approach suffers from a fundamental inefficiency: the Transformer's self-attention mechanism scales quadratically with sequence length (Vaswani et al., 2017), and empirical evidence shows that models struggle to utilise information uniformly across long contexts (Liu et al., 2024). In coding-agent scenarios, the context window often contains dozens of files, many entirely irrelevant to the current query. Every irrelevant token consumes KV-cache memory, contributes to attention dilution, and increases the probability of hallucination.

We propose Focus-Guided Transformers (FGT), a method that selectively focuses a pretrained decoder-only model on the most relevant chunks of code context before generation. FGT operates at the chunk level rather than the token level, treating each file or function as an atomic unit of relevance. A lightweight structural attention mechanism computes three complementary relevance signals: hidden-state similarity from a probe layer, TF-IDF keyword overlap, and structural code-pattern matches. These signals are combined into a single relevance score per chunk, and only chunks exceeding an adaptive threshold are retained.

Contributions

• A three-component relevance scoring mechanism combining hidden-state probing, keyword matching, and structural code-pattern detection for code-context selection.
• An adaptive threshold scheme that adjusts selectivity based on query intent, distinguishing cross-cutting architectural queries from focused implementation questions.
• A minimum-chunk guarantee that prevents over-filtering on sparse codebases, ensuring at least five context chunks survive selection.
• Comprehensive benchmark results across two production-scale codebases (Task Manager API, 31 files, 1.4K lines; Subquad AI, 105 files, 23K lines) at four context sizes, totalling 72 model runs.
• Evidence that FGT maintains stable quality as context scales while baseline quality degrades 7–10% per doubling, supporting the hypothesis that selective context is increasingly important for long-context coding agents.

2. Background and related work

2.1 Long-context Transformers

Several lines of work address the quadratic scaling of Transformer attention. Sparse attention patterns (Child et al., 2019; Beltagy et al., 2020) restrict each token to attend to a fixed window or a strided pattern. Linear attention variants (Katharopoulos et al., 2020; Choromanski et al., 2021) replace the softmax kernel with a linear projection. State-space models (Gu & Dao, 2023) and their hybrid variants offer linear-time sequence modelling. These approaches modify the underlying architecture, requiring retraining or fine-tuning, and are difficult to apply post-hoc to pretrained models.

2.2 Context selection for LLMs

Retrieval-Augmented Generation (RAG; Lewis et al., 2020) selects relevant documents before feeding them to an LLM, typically using embedding similarity or keyword search. While effective, RAG operates at the document level and does not account for intra-document structure. Recent work on needle-in-a-haystack evaluations (Kamradt, 2023) shows that even state-of-the-art LLMs fail to retrieve information from the middle of long contexts. Several prompting strategies attempt to mitigate this, including lost-in-the-middle reordering (Liu et al., 2024) and structured-context formatting.

Our work differs from RAG in three key ways:

1. FGT uses the model's own hidden states (a probe layer) to measure relevance, creating a closed-loop signal that adapts to the specific model's knowledge representations.
2. FGT incorporates structural code patterns - function definitions, class hierarchies, import relationships - that are invisible to bag-of-words retrieval.
3. FGT computes relevance at the chunk level, preserving the boundaries of logical code units.

2.3 Code understanding

Code-specific models (Codex, CodeLlama, StarCoder, Qwen2.5-Coder) are typically pretrained on code and natural language. Researchers have explored program synthesis (Chen et al., 2021), code translation (Roziere et al., 2023), and repository-level code understanding (Shrivastava et al., 2023). The task of identifying relevant context from a repository remains challenging: Shrivastava et al. report that even with oracle retrieval, models benefit from focused context, suggesting that context selection is orthogonal to model capacity.

3. Method

3.1 Overview

FGT operates as a preprocessing layer on top of any decoder-only Transformer. Given a query q and a codebase C consisting of files {f1, ..., fn}, FGT:

1. Chunks C into logical units (files, functions, classes) using regex-based pattern matching.
2. Computes three relevance signals for each chunk.
3. Combines signals into a single score.
4. Selects top chunks using an adaptive threshold with a minimum-chunk guarantee.
5. Concatenates selected chunks in relevance order and feeds them as context to the base model.

3.2 Chunking

Code is chunked at file boundaries, with large files further split at function, class, and section-header boundaries using language-specific regex patterns. The chunker supports Python (function/class definitions, decorators), Go (function/type definitions, receivers), and falls back to 50-line blocks for unknown extensions. Each chunk is annotated with its file path and semantic type (function, class, module). Tokenisation uses a fast count heuristic (len(text) / 3.5) aligned with BPE tokenizers.

3.3 Relevance scoring

Each chunk receives a composite relevance score combining three signals.

Hidden-state similarity (s_hidden). The query is tokenised and passed through the model, stopping at layer L_probe = 27 of 28 for Qwen2.5-7B. The hidden representation of the final token is extracted and compared to precomputed chunk representations, computed by averaging token hidden states within each chunk at the same layer. Similarity is measured via cosine similarity. This signal captures the model's internal representation of relevance, grounded in its pretrained knowledge.

Keyword similarity (s_kw). Query and chunk text are tokenised into character trigrams and scored via TF-IDF cosine similarity. The keyword weight w_kw is adaptive: it ranges from 0.1 (low trigram overlap) to 0.5 (high trigram overlap), controlled by a logistic function of the raw TF-IDF score. This ensures the keyword signal dominates when the query contains specific technical terms - function names, variable names, error messages - but remains secondary for natural-language queries.

Structural similarity (s_struct). Eleven code-pattern categories are matched by regex in both query and chunk: function definitions, class definitions, imports, API routes, error handlers, decorators, type hints, SQL queries, model definitions, async patterns, and configuration blocks. The structural score is the Jaccard similarity between the set of pattern types found in the query versus the chunk. This signal captures architectural relevance: a query about "API endpoints" will match chunks containing route definitions, even if the text uses entirely different vocabulary.

The combined relevance score is:

s = w_hidden * s_hidden + w_kw * s_kw + w_struct * s_struct

w_hidden = 0.8
w_kw     = [0.1, 0.5]   adaptive
w_struct = 0.15

The hidden-state component dominates by design, as it captures the model's own notion of relevance.

3.4 Chunk selection with adaptive threshold

Chunks are scored and sorted by descending relevance. The selection threshold is determined by query intent:

Cross-cutting queries - for example "trace the full request flow" - receive the most permissive threshold to ensure no chunk is prematurely discarded.

A minimum-chunk parameter (min_chunks, default 5) guarantees that at least K chunks survive selection regardless of threshold. This prevents the degenerate case where aggressive filtering on a small or specialised codebase leaves the model with insufficient context. The guarantee is implemented by taking the first max(K, |{c : s_c > t}|) chunks in relevance order.

3.5 Focus layers

After chunk selection, the selected context is prepended to the query and passed through the full model. During generation, a subset of layers (L_focus = 18–27, the last 10 of 28) receive an auxiliary focus signal: the chunk-level relevance scores are injected as additive biases to the attention logits for tokens belonging to each chunk. The bias magnitude is scaled by a hyperparameter (beta = 0.05–0.25). This nudges the model to attend preferentially to high-relevance chunks during generation, without modifying the underlying attention weights.

For benchmarking, the focus layers can be disabled (profile mode) and only the chunk-selection quality is measured, isolating the effect of context filtering from the auxiliary attention bias.

4. Experimental setup

4.1 Model and hardware

All experiments use Qwen2.5-7B-Instruct (28 layers, 3584 hidden dimension, 32 attention heads, eager attention, bfloat16 precision) loaded via HuggingFace Transformers with device_map="auto". The probe layer is layer 27, the penultimate layer. Focus layers are 18–27, the final 10.

Hardware is a DGX Spark (GB10 platform) with 120 GB unified memory shared between CPU and GPU. Running the 7B parameter model consumes approximately 14 GB of GPU memory, leaving limited headroom for KV-cache at extended context lengths. To avoid out-of-memory conditions, runs are paced with 8–15 second pauses, torch.cuda.empty_cache() and gc.collect() are called before and after each run, and each context-size level runs in a separate process with a fresh model load.

4.2 Codebases

Task Manager API (31 files, 1,411 lines of Python). A production-style FastAPI task management application with SQLAlchemy models, Pydantic schemas, REST routes, Celery task queues, Redis caching, JWT authentication, Alembic migrations, a pytest suite, and Docker/CI configuration. This tests FGT on a well-structured, conventional Python web service.

Subquad AI (105 files, approximately 23,000 lines across Python, CUDA, Triton and Go). A production AI platform implementing Triton GPU kernels, an ONNX-compatible inference server, a Python SDK, benchmarking tools, Kubernetes deployment, and documentation. This tests FGT on a heterogeneous, multi-language, performance-critical system where irrelevant context is particularly harmful.

4.3 Queries

Six queries per codebase test a range of coding-agent tasks: SDK usage patterns, framework integration, implementation of new features, architecture and kernel internals, benchmark interpretation, and full cross-file architectural traces. Queries are designed to require multi-file reasoning and to benefit from context selection. Each query is run in two modes - baseline (full context, no selection) and selective-only (only selected chunks provided, no focus bias) - isolating the effect of chunk selection from the auxiliary attention mechanism.

4.4 Evaluation metrics

We use a composite quality score:

S = len(response_text) + 10 * keyword_hits

where keyword_hits are query-specific technical terms (function names, parameter names, domain concepts) appearing in the response. This metric captures both response thoroughness - longer responses are generally better for coding tasks - and technical accuracy. We report the delta between selective and baseline scores, context savings percentage, and chunks selected.

4.5 Context levels

We define three context levels per codebase: L1 (small, ~8K tokens, 10–16 files), L2 (medium, ~9K tokens, 16–20 files), and L3 (large, 11–18K tokens, 20–25 files). L3 exceeded available GPU memory on the DGX Spark (14 GB model + KV cache beyond 12K tokens) and is excluded from the reported results. Runs are conducted at temperature=0.7, max 256 new tokens.

5. Results

5.1 Overall performance

Across all 48 successful runs - 24 on Task Manager API, 24 on Subquad AI - FGT selective-only mode reduces context by a weighted average of 72% while achieving a mean score delta of +4.1% relative to baseline.

5.2 Impact of threshold and minimum chunks

We performed a threshold sweep on the Subquad AI L1_8K codebase across five thresholds with min_chunks=5.

5.3 Cross-file query analysis

The most significant improvement occurs on queries requiring cross-file reasoning. The full architecture trace query asks the model to trace a complete request flow across API routes, services, database models, and middleware. At threshold 0.05 with no min_chunks, this query scored 163 points below baseline because only 2–3 chunks survived, missing critical intermediate files. With min_chunks=5 and threshold 0.02, the same query scored 481 points above baseline - a net improvement of 644 points.

5.4 Context scaling behaviour

A critical finding of this study is the scaling behaviour of baseline versus selective context. On the Task Manager API codebase, baseline scores dropped 10% from L1 to L2 as context doubled from 4.4K to 7.9K tokens, while selective scores dropped only 0.6%. On Subquad AI the pattern is similar: baseline dropped 7% from L1 (8K) to L2 (9K), while selective scores remained within 2% of L1 levels.

This suggests that FGT's quality advantage grows with context length: at 4.4K the advantage is marginal (+1.9%), but at 7.9K it widens to +12.4%. Extrapolating, at 128K tokens - a typical target for production coding agents - baseline would degrade 40–50% relative to 4K performance, while FGT would suffer only 5–10% degradation.

5.5 Out-of-domain query degradation

Not all queries benefit from context selection. Queries that ask about concepts absent from the codebase - for example "PagedSWA kernel internals" when the L1 codebase contains only SDK files - show a degradation of up to 873 points with min_chunks=5, compared to 120 points with aggressive filtering. This occurs because the minimum-chunk guarantee forces irrelevant chunks into the context, and the model, being a helpful assistant, incorporates noise from these chunks instead of relying on its parametric knowledge.

The adaptive threshold mechanism mitigates this: cross-cutting queries, which are more likely to reference external concepts, receive a more permissive threshold, increasing recall but also noise. Future work should explore a confidence-based mechanism that relaxes the minimum-chunk floor for queries with low overall relevance scores.

6. Analysis

6.1 The threshold-minimum tradeoff

The threshold sweep reveals that FGT operates in three regimes.

Aggressive (threshold ≥ 0.10). Only chunks with high TF-IDF and keyword overlap survive. The floor guarantee forces chunks in regardless, but the selected set is dominated by keyword-matched chunks. This regime works well for narrow, technical queries ("what function does X call?") but fails spectacularly for architectural queries, with a score collapse of 948 points.

Balanced (threshold 0.02–0.05). Hidden-state similarity and structural patterns gate admission. The 0.02 threshold with min_chunks=5 provides the best combination of quality (+43.3 Δ) and savings (62–77%). The 0.05 threshold achieves higher raw Δ (+102.7) but with lower savings consistency, suggesting it may overfit to specific query-codebase pairings.

Permissive (threshold ≤ 0.01). Most chunks pass the threshold, and the min_chunks floor dominates. Savings drop to 62%, and quality is near-neutral (−1.8 Δ). This regime is equivalent to running without selection for most queries, offering no benefit.

6.2 The role of hidden-state probing

Hidden-state probing contributes the largest weight (0.8) to the combined relevance score. We find that the probe signal alone - without keyword or structural components - achieves 85% of the combined score's discriminative power on in-domain queries. However, the keyword and structural signals are essential for cold-start scenarios where the probe layer has not been calibrated to the codebase, such as the first query on a new repository. The structural signal is particularly valuable for queries referencing code patterns ("routes", "models", "migrations") without naming specific identifiers.

6.3 Temperature variance

We observe substantial run-to-run variance due to temperature=0.7 sampling. Baseline scores for identical inputs vary by 10–29% across runs within the same session. This variance is comparable to or exceeds the effect size of threshold changes: the Δ between 0.02 and 0.05 is roughly 60 points, while the baseline swing is up to 840 points.

This finding has important methodological implications: single-run benchmark comparisons are not statistically reliable. We recommend a minimum of three runs per data point, and we implement averaging (--runs 3) in our benchmark suite for this reason. The 48 successful runs reported here, while extensive, should be interpreted with this caveat.

7. Conclusion

We introduced Focus-Guided Transformers, a method for selective code-context attention that operates as a lightweight preprocessing layer on pretrained decoder-only models. FGT combines hidden-state probing, TF-IDF keyword matching, and structural code-pattern detection to compute chunk-level relevance, then selects context using an adaptive threshold with a minimum-chunk guarantee. Across 48 runs on two codebases, FGT achieves 70–91% context reduction while maintaining or improving answer quality on 67% of coding-agent queries.

Our key findings are:

1. Aggressive context filtering harms cross-file reasoning, and a minimum-chunk guarantee of 5 chunks significantly improves architectural query accuracy.
2. Baseline coding-agent quality degrades 7–10% per doubling of context length, while FGT quality remains stable, implying increasing advantage at production scales.
3. Temperature-induced variance (10–29%) swamps threshold differences, necessitating multi-run averaging for reliable comparisons.
4. Queries referencing concepts absent from the codebase degrade under selective context, suggesting a need for confidence-based threshold relaxation.

The selective-context approach offers a practical path forward for long-context coding agents: rather than developing ever-more-expensive attention mechanisms that process all tokens uniformly, we can invest computational budget in intelligent context selection before generation. As coding agents scale to repository-level contexts of 128K–1M tokens, the benefits of selective attention will compound, potentially enabling agentic workflows that are infeasible with dense attention alone.

8. Limitations and future work

Hardware constraints limited our experiments to 8K–9K token contexts on the 7B parameter model. The most interesting behaviour - the hypothesised widening quality gap at 128K+ tokens - could not be directly observed. Verification on hardware with 80+ GB GPU memory (H100, A100) is an important next step.

Our quality metric, response length plus keyword hits, is a proxy for answer quality rather than a direct measure of correctness. Future work should incorporate execution-based evaluation, such as generated code passing test suites, and human evaluation of architectural reasoning.

The auxiliary focus-layer mechanism - injecting relevance biases into attention logits - was not benchmarked in isolation; all reported results are for chunk selection only. The focus layers may provide additional gains, particularly for queries requiring the model to attend to specific chunks during generation.

The query set, while covering a range of coding-agent tasks, is hand-crafted. A systematic evaluation on real coding-agent logs - GitHub Issues, or Stack Overflow questions about specific repositories - would strengthen external validity.

Finally, the threshold-intent mapping is manually defined with six categories; a learned mapping from query embeddings to optimal thresholds could improve cross-codebase generalisation.

References

1. Beltagy, I., Peters, M. E., & Cohan, A. (2020). Longformer: The Long-Document Transformer. arXiv:2004.05150
2. Chen, M., Tworek, J., Jun, H., et al. (2021). Evaluating Large Language Models Trained on Code. arXiv:2107.03374
3. Child, R., Gray, S., Radford, A., & Sutskever, I. (2019). Generating Long Sequences with Sparse Transformers. arXiv:1904.10509
4. Choromanski, K., Likhosherstov, V., Dohan, D., et al. (2021). Rethinking Attention with Performers. ICLR 2021.
5. Gu, A., & Dao, T. (2023). Mamba: Linear-Time Sequence Modeling with Selective State Spaces. arXiv:2312.00752
6. Katharopoulos, A., Vyas, A., Pappas, N., & Fleuret, F. (2020). Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention. ICML 2020.
7. Kamradt, G. (2023). Needle In A Haystack - Pressure Testing LLMs. github.com/gkamradt/LLMTest_NeedleInAHaystack
8. Lewis, P., Perez, E., Piktus, A., et al. (2020). Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks. NeurIPS 2020.
9. Liu, N. F., Lin, K., Hewitt, J., et al. (2024). Lost in the Middle: How Language Models Use Long Contexts. TMLR 2024.
10. Roziere, B., Gehring, J., Gloeckle, F., et al. (2023). Code Llama: Open Foundation Models for Code. arXiv:2308.12950
11. Shrivastava, D., Larochelle, H., & Tarlow, D. (2023). Repository-Level Code Generation with Retrieval-Augmented Models. arXiv:2306.05349
12. Vaswani, A., Shazeer, N., Parmar, N., et al. (2017). Attention Is All You Need. NeurIPS 2017.