CREME: Robustness Enhancement of Code LLMs via Layer-Aware Model Editing

TL;DR

CREME is a novel method that enhances code LLM robustness by identifying sensitive layers and applying targeted parameter edits, significantly improving performance on perturbed prompts while preserving accuracy on clean inputs.

Contribution

CREME introduces a layer-aware model editing technique that improves LLM robustness to prompt perturbations by selectively updating parameters in sensitive layers.

Findings

CREME improves Pass@1 accuracy by 63% on perturbed prompts.

Robustness-sensitive layers are mainly in middle and deeper layers.

Locations of sensitive layers vary across different model architectures.

Abstract

Large language models (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.