Multi-task LLMs for Bug Classification: Efficient Inference with Auxiliary Decoding Heads

Abstract: The rapid adoption of LLM-powered code generation has dramatically accelerated software development, yet effective verification methods remain severely underdeveloped. Existing bug localization techniques are either prohibitively expensive, requiring minutes of agentic reasoning and thousands of generated tokens per file, and/or operate at coarse function-level granularity unsuitable for precise debugging. While works that focus on line-level granularity and are more light-weight are often limited in their performance or context size. We introduce a novel line-level bug localization approach that addresses these limitations through three key contributions: (1) a token alignment algorithm that overcomes fundamental tokenization challenges in previous work, (2) a lightweight multi-task LLM for bug localization (MLC) enabling efficient line-level bug classification, and (3) an optimized training recipe for multi-line prediction. Our method achieves state-of-the-art performance among similar setups on line-level bug localization with full-file context. At the same time we reach comparable performance to agentic approaches on Defects4J and PypiBugs benchmarks while reducing inference latency by orders of magnitudes, requiring only a single generated token per file. We further demonstrate strong generalization by introducing and evaluating on a small out-of-domain evaluation datasets in Python. We will open source our code, models, and datasets upon acceptance.

• The paper introduces an auxiliary decoding head that efficiently enables line-level bug classification, addressing token-line misalignment.
• It employs a robust token-line alignment algorithm and multi-task aggregation strategies to optimize bug indicator prediction.
• Experimental results show significant improvements in Top1 and Top5 accuracy with faster inference compared to traditional methods.

Multi-task LLMs for Bug Classification: Advancing Efficient Line-level Bug Localization via Auxiliary Decoding Heads

Introduction and Motivation

The integration of LLMs in code generation pipelines has considerably accelerated software development, while effective automated verification and bug localization methods remain a major bottleneck. Existing techniques for fault localization—either agentic LLM architectures or fine-tuned discriminative models—suffer from high inference latency, large computational costs, and inadequate granularity (often function-level rather than line-level). Moreover, tokenization challenges in standard LLM workflows often lead to imprecise bug attribution at the line level, as token boundaries rarely align with code line breaks. This paper introduces a multi-task LLM approach, termed MLC, which overcomes these limitations through a direct, efficient auxiliary bug-classification head, optimized token-line alignment, and enhanced training paradigms for multi-line prediction.

Figure 1: Model architecture showing both the original LLM ‘channel’ for next token prediction as well as the additional binary bug classification head.

Methodology

The proposed approach delivers efficient, precise bug localization via three components: (i) a token alignment algorithm resolving misalignment in standard tokenization strategies, (ii) an auxiliary decoding head for multi-task, line-level bug classification, and (iii) a cost-effective training regime that enables single-pass classification over entire files.

Token-Line Alignment

Tokenization in modern LLMs (e.g., GPT-5) does not preserve code line boundaries due to the nature of subword encodings, with newline indicators often split across or conflated with adjacent tokens.

Figure 2: Code snippet illustrating GPT-5 tokenizer splitting the newline token across multiple lines and tokens, complicating direct mapping between code lines and tokens.

The paper introduces a robust alignment algorithm employing a two-pointer heuristic to map tokens to code lines with minimal error, ensuring that bug labels are attributed accurately at the line granularity.

Multi-task Decoding Architecture

By attaching an additional binary classification head to the frozen backbone of a pre-trained code-oriented LLM (e.g., Qwen2.5-Coder or CodeGen), line representations are aggregated via mean, sum, or last-token operations over tokens corresponding to each line. A lightweight feedforward network projects these line embeddings into bug/non-bug logits. Crucially, the auxiliary head does not disturb the original next-token prediction pathway, thus maintaining the base LLM’s general code modeling capacity while affording simultaneous bug classification.

Figure 3: Full inference flow from a code file to per-line bug probability outputs, leveraging the alignment algorithm and auxiliary decoding head for efficient classification.

Training and Inference Paradigm

MLC is trained using a weighted binary cross-entropy objective over all lines in a code file (batch size = full file), addressing severe data imbalance typical in real-world datasets (buggy lines are rare). The approach supports both PEFT (e.g., LoRA adapters) and head-only fine-tuning, allowing flexible trade-offs between efficiency and adaptation. Unlike generative or agentic LLM-based approaches, inference requires only a single pass to yield bug probabilities for all lines, reducing end-to-end latency by up to several orders of magnitude.

Experimental Evaluation

Benchmarks and Dataset Coverage

Evaluation is conducted on Defects4J (Java), PypiBugs (Python), and a newly proposed out-of-domain Python dataset (BugsEval), with comprehensive splits and exhaustive annotation. The MLC technique is assessed against both traditional static/dynamic analysis baselines (Ochiai, DeepFL) and recent state-of-the-art LLM-based methods (LLMAO, FlexFL, MemFL, AgentFL).

Quantitative Results

On Defects4J, MLC with the PEFT Qwen1.7B achieves 16.3% Top1 and 39.5% Top5 line-level accuracy—well above classical methods and direct SFT baselines, with SFT yielding only 0% Top1 and 14.3% Top5. Notably, without PEFT adaptation, classification head fine-tuning on Defects4J remains competitive.

On PypiBugs, MLC consistently outperforms SFT and shows that smaller Qwen3 models match or surpass SFT on both Top1 and Top5.

In comparison to short-context, function-level agentic approaches (e.g., MemFL, AgentFL), the largest MLC model (CodeGen 16B) yields 28.6% Top1 and 71.4% Top5, rivaling the best agentic Top5 results but with dramatically faster inference and lower resource requirements.

Empirical ablations indicate that the last-token as an aggregation function and sigmoid activation in the decoding head marginally outperform softmax/mean aggregations, highlighting subtle architectural biases for causal LLM heads.

Analysis: Generalization and Hyper-Parameter Effects

Testing on the out-of-domain BugsEval (Python) dataset after PypiBugs training reveals that Top1 generalization sometimes exceeds in-distribution performance (e.g., 13.8% Top1 for Qwen8B), while Top5 drops, suggesting the auxiliary head captures general defectiveness cues but struggles with subtle, domain-specific bug reasoning in OOD settings. This demonstrates that fundamental bug localization abstractions learned by MLC generalize, though identifying all buggy lines in OOD cases remains challenging.

Hyper-parameter examination reveals that the base LLM already provides strong code representations for defect detection. Additional decoding layers do not guarantee improvements due to already rich feature extraction; aggregation strategies and dynamic loss weighting further modulate performance.

Practical and Theoretical Implications

MLC’s architectural efficiency—requiring a single forward pass per file for line-level predictions—enables practical deployment for large-scale, continuous integration systems previously constrained by the prohibitive costs of agentic LLM approaches. The modularity of the aggregation strategy and auxiliary head introduces an extensible framework for multi-task LLMs, pointing toward a future where LLMs routinely handle discrete, independent tasks via dedicated output channels with minimal interference.

Theoretically, the findings contest the prevailing assumption that SFT with prompt engineering suffices for code bug localization. Instead, explicit multi-task channelization and careful token-level alignment are necessary for robust and efficient addressal of fine-grained debugging tasks.

Future Directions

Several research trajectories are highlighted:

• Extending the token alignment algorithm to function or project granularity to bridge gaps with function-level agentic baselines.
• Scaling multi-channel paradigms to additional tasks (e.g., style, security bugs) or output modalities (e.g., span annotation).
• Leveraging stronger LLMs to further increase localization accuracy without incurring agentic inference costs.
• Systematic study of multi-channel interference and transfer in shared LLM backbones—crucial for ‘poly-task’ coding assistants.

Conclusion

This work introduces a multi-task LLM architecture for line-level bug localization that approaches or surpasses existing state-of-the-art in both agentic and discriminative LLM regimes, while achieving substantive gains in computational efficiency. The combination of a robust token alignment mechanism and an efficient auxiliary classification head enables competitive, fast, and flexible bug localization—supporting both practical deployment and foundational LLM research into task specialization.

This paradigm emphasizes the need for architectural multi-tasking and precise alignment algorithms in effective code intelligence with LLMs, indicating a near-term future in which modular, efficient, multi-channel LLMs replace monolithic, sequential reasoning systems for granular code analysis.