Breaking Memorization Barriers in LLM Code Fine-Tuning via Information Bottleneck for Improved Generalization

TL;DR

The paper identifies a memorization barrier in LLM code fine-tuning, where pre-memorized downstream data in the base model hampers effective adaptation. It proposes IB-regularized fine-tuning (IB-FT), which uses an information bottleneck to compress spurious memorized features while preserving task-relevant information, formulated via $I(X;Z)$ and $I(Z;Y)$ terms and implemented with a variational encoder. Across OriGen and Evol-CodeAlpaca-V1, IB-FT yields stronger and more stable generalization than conventional fine-tuning, with large gains at Pass@$1$ and improved robustness under strict multi-sample criteria such as Pass@$k^{(m)}$ and varying decoding temperatures. The work provides a practical, principled approach to breaking memorization-driven limits in code generation, with implications for safer, more reliable code synthesis and broader applicability to domain-adaptive LLM fine-tuning.

Abstract

Adapting pretrained large language models (LLMs) to code domains via supervised fine-tuning (FT) has been commonly used for code generation. However, we identify a previously underappreciated failure mode, the memorization barrier, where strong memorization of downstream code data in the base model could trap optimization and prevent the standard FT from effectively acquiring new, generalizable code knowledge. To overcome this barrier, we propose the information bottleneck (IB)-guided fine-tuning, termed IB-FT, which applies an IB penalty on hidden representations of the code data to compress spurious, memorized features while preserving task-relevant information. Extensive experiments on two code benchmarks (OriGen and Evol-CodeAlpaca-V1) show that IB-FT substantially alleviates the memorization barrier, improves top-1 performance (Pass@$1$), and yields far more stable gains under the stricter multi-sample metric Pass@$k^{(m)}$ (a problem counts as solved only if at least $m$ of $k$ samples pass unit tests) compared with conventional FT.

Figures (7)

• Figure 1: Schematic overview of IB-FT versus conventional FT and a representative performance comparison when fine-tuning on OriGen. (a) A code fine-tuning "race": conventional FT (left) appears to advance quickly by memorization but follows a winding, less-generalizable path, whereas IB-FT (right) progresses along a straighter trajectory toward the "winning" generalization. (b) Test-time performance for DeepSeek-Coder-7B-Instruct-v1.5, reported on Eval-Human and Eval-Machine using Pass@$k^{(m)}$ with $m=k\in\{1,5,10\}$ (a problem is counted only if all $k$ samples pass). IB-FT consistently attains higher, more stable accuracy and occupies the high-performance region compared to FT.
• Figure 2: Performance of DeepSeek-Coder-7B-Instruct-v1.5 fine-tuned on the OriGen dataset, evaluated on the Eval-Human and Eval-Machine. (a) Pass@$k$ results for $k \in \{10, 5, 1\}$, measuring the probability that at least one of the top-$k$ generations passes the test cases. (b) Pass@$k^{(m)}$ results with $k=10$ and $m \in \{2,5,10\}$, requiring multiple successful generations among the ten samples to reflect robustness.
• Figure 3: Downstream data already memorized during pretraining. Memorization patterns measured by Min-K% Prob ($K=20$) using DeepSeek-Coder-7b-Instruct-v1.5 and Llama-3 8B as base models on two code datasets, OriGen and Evol-CodeAlpaca-V1. The x-axis shows Min-20% Prob scores and the y-axis their density. For each dataset, three distributions are reported: $p_{0}(\mathcal{D}_\mathrm{code}; {\boldsymbol{\theta}}_0)$ (code dataset on the base model), $p_{1}(\mathcal{D}_\mathrm{code}; {\boldsymbol{\theta}}_\mathrm{code})$ (code dataset on the fine-tuned model), and $p_{2}(\mathcal{D}_\mathrm{TOFU}; {\boldsymbol{\theta}}_0)$ (TOFU dataset on the base model), where $p_2$ provides the weakest-memorization reference.
• Figure 4: Pass@$k$ performance of the fine-tuned model on the OriGen dataset using DeepSeek-Coder-7b-Instruct-v1.5, evaluated by (a) Eval-Human and (b) Eval-Machineliu2023verilogeval, with different pruning ratios of highly memorized fine-tuning data identified in Fig. \ref{['fig:mink_dis']}. The x-axis shows the removal ratio of the most memorized examples (0% corresponds to full dataset), and the y-axis reports Pass@$k$ for $k \in \{1,5,10\}$.
• Figure 5: Representation distance and angle analysis with fine-tuning data on OriGen, using the base DeepSeek-Coder-7B-Instruct-v1.5 model and its fine-tuned variants (FT and IB-FT). Hidden representations $h_{\boldsymbol{\theta}}(x)$ are extracted from the 20th layer, and paired samples are drawn between the most-memorized and least-memorized groups (identified by Min-20% Prob scores in Fig. \ref{['fig:mink_dis']}) to compute (left) their $\ell_2$ distance distributions and (right) angular disparity distributions. Dashed vertical lines mark distinct separation patterns under different fine-tuning regimes.