Improving Code LLM Robustness to Prompt Perturbations via Layer-Aware Model Editing

Abstract: LLMs have demonstrated impressive capabilities in code generation, where the natural language prompt plays a crucial role in conveying user intent to the model. However, prior studies have shown that LLMs are highly sensitive to prompt perturbations. Minor modifications in wording, syntax, or formatting can significantly reduce the functional correctness of generated code. As perturbations frequently occur in real-world scenarios, improving the robustness of LLMs to prompt perturbations is essential for ensuring reliable performance in practical code generation. In this paper, we introduce CREME (Code Robustness Enhancement via Model Editing), a novel approach that enhances LLM robustness through targeted parameter updates. CREME first identifies robustness-sensitive layers by comparing hidden states between an original prompt and its perturbed variant. Then, it performs lightweight parameter editing at the identified layer to reduce performance degradation. We evaluate CREME on two widely used code generation benchmarks (HumanEval and MBPP) along with their perturbed counterparts. Experimental results show that CREME improves Pass@1 accuracy by 63% on perturbed prompts while maintaining stable performance on clean inputs, with accuracy deviations within 1%. Further analysis reveals that robustness-sensitive layers are primarily concentrated in the middle and deeper layers of the network, and their locations vary across different model architectures. These insights provide a valuable foundation for developing future robustness-oriented editing strategies.

• The paper introduces CREME, which uses layer-aware editing to significantly enhance code LLM robustness against minor prompt changes.
• It leverages causal tracing to identify key layers and applies targeted updates that recover up to a 63% improvement in Pass@1 on perturbed prompts.
• The approach preserves clean prompt performance within a ±1% margin and generalizes across 20 distinct perturbation types.

Improving Code LLM Robustness to Prompt Perturbations via Layer-Aware Model Editing

Introduction and Motivation

The paper addresses a critical limitation in code-oriented LLMs: their pronounced sensitivity to minor prompt perturbations. Empirical evidence demonstrates that even semantically negligible changes—such as typographical errors, redundant words, or slight rephrasings—can cause substantial degradation in the functional correctness of generated code. This fragility is especially problematic in real-world software engineering workflows, where prompt variability is the norm and not all users are expert prompt engineers.

Existing robustness enhancement strategies fall into two main categories: input-level interventions (e.g., prompt denoising, paraphrasing) and model-augmentation approaches (e.g., soft prompts, LoRA, RAG). However, these methods either do not alter the model’s inherent robustness or introduce additional system complexity and retraining requirements. The paper proposes a lightweight, model-internal solution: CREME (Code Robustness Enhancement via Model Editing), which leverages targeted parameter updates to improve robustness to prompt perturbations without full retraining or architectural changes.

CREME: Methodology

Task Formulation

Given a pre-trained code LLM $G: P \rightarrow Y$ mapping a prompt $P$ to code $Y$, and a pair of prompts $(P_{ori}, P_{pert})$ that are semantically equivalent but differ in surface form, the goal is to construct an edited model $G_{\mathcal{W}'}$ such that $G_{\mathcal{W}'}(P_{pert}) \approx G(P_{ori})$, while preserving performance on clean prompts.

Key Layer Localization via Causal Tracing

CREME introduces a layer-wise causal intervention to identify the "robustness-sensitive" layer(s) responsible for the model’s output divergence under perturbation. The process is as follows:

1. Baseline Evaluation: Compute pass@1 accuracy for $P_{ori}$ and $P_{pert}$.
2. Layer-wise Intervention: For each transformer layer $L_\ell$, replace the hidden states of $P_{pert}$ at $P$0 with those from $P$1, then propagate forward and evaluate the output.
3. Restoration Improvement Metric: For each layer, compute the fraction of the accuracy gap between $P$2 and $P$3 that is recovered by the intervention. The layer with maximal restoration is selected as the key layer $P$4.

This approach is distinct from prior knowledge editing methods, which typically rely on subject tokens or fixed output references, and is tailored for the more complex, multi-sentence prompts in code generation.

Targeted Parameter Editing

Once $P$5 is identified, CREME performs a localized parameter update on the MLP output projection matrix $P$6 at that layer. The update is guided by two objectives:

• Alignment Loss: Minimize the MSE between the hidden states of $P$7 and $P$8 at $P$9.
• Preservation Loss: Penalize deviation of the hidden state for $Y$0 after editing, ensuring the model’s behavior on clean inputs is retained.

The total loss is:

$Y$1

where $Y$2 is a regularization coefficient.

Only $Y$3 is updated via gradient descent, with all other parameters frozen. Early stopping is employed to prevent overfitting.

Experimental Evaluation

Datasets and Perturbations

Experiments are conducted on HumanEval and MBPP, two standard Python code generation benchmarks, and their perturbed variants generated using NLPerturbator. Perturbations span 20 types, including lexical (typos, deletions, insertions), syntactic (paraphrasing, word swaps), and combinations thereof. Only prompts with a robustness drop $Y$4 are retained for evaluation.

Baselines

CREME is compared against:

• Self-Denoising: Prompt-level denoising using the LLM itself.
• LoRA Fine-tuning: Representation alignment via LoRA modules.
• ROME: Causal tracing-based factual knowledge editing.
• DINM: Detoxification via knowledge editing with gold-standard responses.

Metrics

• Pass@1: Fraction of generations passing all test cases on first attempt.
• G-RIR (Generalized Relative Improvement Ratio): Measures generalization of robustness improvement across unseen prompts of the same perturbation type.

Results

Effectiveness

CREME achieves a 63% average improvement in Pass@1 on perturbed prompts, outperforming all baselines. The average G-RIR is 0.37, a 41% gain over the best baseline (LoRA). Notably, CREME maintains performance on clean prompts within a $Y$51% margin, indicating minimal negative side effects.

Generality

CREME generalizes across all 20 perturbation types, with higher G-RIR for word-level perturbations (e.g., typos, deletions) and lower for sentence-level rephrasings. Performance is stronger on MBPP (shorter, more templated prompts) than on HumanEval (longer, more diverse prompts).

Ablation Study

Removing key components (layer localization, early stopping, preservation loss) degrades G-RIR by 12–20%, with layer localization being most critical. This underscores the necessity of precise causal tracing for effective robustness enhancement.

Layer Distribution Analysis

Robustness-sensitive layers are not uniformly distributed; they cluster in the middle-to-deep layers, with their locations varying by model architecture and perturbation type. For example, CodeLlama’s key layers are often in the last few layers, while Qwen2.5-Coder’s are more centralized. Editing perturbations tend to shift the key layer deeper, reflecting increased semantic complexity.

Implementation Considerations

Practical Workflow

1. Collect a pair of prompts $Y$6 for a given perturbation type.
2. Run causal tracing to localize the key layer $Y$7.
3. Edit $Y$8 using the alignment and preservation losses, with early stopping.
4. Deploy the edited model for improved robustness to similar perturbations.

Resource Requirements

• Editing is lightweight: only a single layer’s parameters are updated, and the process converges in $Y$9 steps.
• No full retraining or additional data is required.
• The method is compatible with standard PyTorch/HuggingFace LLM implementations.

Limitations

• Generalization is strongest within the same perturbation type; cross-type robustness is not guaranteed.
• The approach is evaluated on open-source models; applicability to proprietary LLMs (e.g., GPT-4) is untested.
• The method assumes access to model internals for parameter editing, which may not be feasible in all deployment scenarios.

Theoretical and Practical Implications

CREME demonstrates that localized, layer-aware model editing can substantially improve LLM robustness to prompt perturbations with minimal side effects. The findings suggest that robustness is governed by specific subregions of the model, and that targeted interventions can yield broad improvements within perturbation categories. This challenges the prevailing reliance on input-level or global model modifications for robustness.

Theoretically, the work extends the scope of knowledge editing from factual updates to robustness enhancement, highlighting the causal role of intermediate representations in mediating model sensitivity. Practically, CREME offers a scalable, data-efficient solution for post-hoc robustness tuning in code LLMs, with potential applicability to other domains (e.g., QA, summarization) pending further research.

Future Directions

• Cross-domain generalization: Extending CREME to non-code LLM tasks and more adversarial perturbations.
• Multi-layer or multi-point editing: Investigating whether editing multiple layers can further enhance robustness.
• Automated perturbation detection: Integrating CREME with systems that automatically detect and categorize prompt perturbations in real time.
• Closed-source LLMs: Exploring the feasibility of similar interventions via API-based or black-box model editing.

Conclusion

CREME provides a principled, efficient framework for enhancing code LLM robustness to prompt perturbations via layer-aware model editing. By leveraging causal tracing and targeted parameter updates, it achieves substantial improvements in functional correctness on perturbed prompts, with minimal impact on clean data. The approach offers actionable insights for both the mechanistic understanding of LLM robustness and the practical deployment of more reliable code generation systems.