SemRep: Generative Code Representation Learning with Code Transformations

1. Problem Statement

Code transformation is a critical task in software engineering, ranging from bug fixing and refactoring to performance optimization (e.g., GPU kernel tuning). Existing approaches typically treat code transformation as an end-to-end learning task, where Large Language Models (LLMs) implicitly learn the necessary semantic representations within their weights.

Key Limitations of Current Methods:

• Implicit Semantics: Models rely on latent weights to understand code semantics, which often leads to errors when the model fails to disentangle what must be preserved (original functionality) from what can be changed (optimization).
• Rigid Abstractions: Alternative approaches using compiler-level abstractions (e.g., ASTs, Control Flow Graphs) are often too rigid or expensive to construct on the fly for natural language-instructed editing.
• Lack of Exploration: Direct transformation often conflates rewriting and optimization in a single step, causing the model to make incorrect edits or miss optimization opportunities (e.g., failing to fuse operations or optimize memory access).

2. Methodology: The SEMREP Framework

SEMREP introduces a two-stage framework that explicitly decouples semantic understanding from instruction-specific transformation. The core insight is to use semantics-preserving transformations as an intermediate, generative representation.

A. Core Concept: Generative Code Representation Learning

Instead of treating code representation as a latent vector, SEMREP trains the model to generate explicit, semantically equivalent source code ($C_{rep}$) as an intermediate step.

• Definition: Two programs are equivalent if they produce identical outputs for a given set of test cases ($U_{sem}$).
• Goal: The model learns to explore the "semantic equivalence class" of the input code, generating diverse variants (e.g., different loop structures, API calls) that maintain functionality but alter the internal structure.

B. Training Paradigm (Two Stages)

The framework utilizes Group Relative Policy Optimization (GRPO), a reinforcement learning (RL) algorithm, with verifiable rewards based on test execution.

1. Stage 1: Generative Code Representation Learning (Mid-Training)

   • Objective: Train the model to generate semantically equivalent variants of the source code ($C_{src} \to C_{rep}$).
   • Reward Function ($R_{sem}$): Rewards compilability and semantic invariance (passing $U_{sem}$ tests). It encourages diversity in syntax and structure to foster creative downstream transformations.
   • Budget: Uses a portion ($\alpha$) of the total training budget.

2. Stage 2: Instruction-Specific Transformation Learning (Fine-tuning)

   • Objective: Train the model to transform the intermediate representation ($C_{rep}$) into the final target code ($C_{tgt}$) that adheres to user instructions (e.g., "optimize for speed").
   • Reward Function ($R_{inst}$): Rewards compilability, semantic invariance (ensuring $C_{tgt}$ still passes $U_{sem}$), and instruction adherence (passing $U_{edit}$ tests).
   • Budget: Uses the remaining portion ($1-\alpha$) of the training budget.

Note: The total training budget for SEMREP is strictly matched to baseline fine-tuning models to ensure fair comparison.

C. Inference: Test-Time Scaling via Evolutionary Search

SEMREP supports iterative inference where the model alternates between:

1. Semantic Exploration: Generating diverse semantically equivalent variants.
2. Instruction-Specific Transformation: Applying specific edits to the best candidates.
   This process creates a beam of candidate programs, allowing the system to scale compute at inference time to find optimal solutions that a single-pass generation might miss.

3. Key Contributions

1. Explicit Semantic Representation: Proposes generating semantically equivalent code as an explicit intermediate representation, making reasoning about code behavior verifiable and interpretable.
2. Decoupled Training Strategy: Separates the learning of "what the code does" (semantics) from "how to change it" (instruction following), preventing the model from conflating optimization with functional correctness.
3. Efficiency and Generalization: Demonstrates that smaller models (e.g., 32B parameters) trained with SEMREP can outperform significantly larger models (e.g., 685B) and closed-weight commercial models.
4. Synergy with Evolutionary Search: The explicit intermediate representations make the model highly amenable to evolutionary coding agents, enabling the discovery of complex optimizations (like algorithmic shifts) that greedy search methods miss.

4. Experimental Results

The authors evaluated SEMREP on EditBench (general code editing) and KernelBench (GPU kernel optimization).

• Performance vs. Baselines:
  • Correctness: Outperformed extensively fine-tuned baselines by 6.9% on EditBench.
  • Performance: Achieved 1.1× speedup on KernelBench compared to baselines with the same training budget.
  • Generalization: Improved cross-hardware generalization by 13.9% (tested on unseen H200 GPUs) and robustness to code perturbations by 6.7%.
• Model Size Efficiency: A 32B model (QwQ-32B) with SEMREP matched or exceeded the performance of 12× larger models (e.g., 405B Qwen3-Coder) and commercial models like GPT-o4-mini.
• Evolutionary Agent Integration: When combined with an evolutionary agent (OpenEvolve), SEMREP discovered optimizations for Minimum Spanning Tree and Circle Packing problems that 685B models failed to find, achieving state-of-the-art results with 25% less inference compute.
• Case Studies:
  • CUDA Kernels: Successfully fused operations and optimized memory access (tiling, shared memory) where direct transformation methods failed.
  • Parallelism: In a case study involving asyncio, SEMREP correctly separated refactoring (extracting helper functions) from parallelization, avoiding the output-ordering bugs common in direct transformation approaches.

5. Significance and Impact

• Paradigm Shift: Moves code editing from "black-box" end-to-end learning to a structured, verifiable process where semantic understanding is an explicit, generative capability.
• Scalability: Proves that explicit representation learning allows smaller models to compete with massive models, reducing the barrier to entry for high-quality code generation.
• Reliability: By enforcing semantic equivalence via test execution during training and inference, SEMREP significantly reduces the risk of introducing functional bugs during optimization.
• Future Directions: The framework opens new avenues for combining LLMs with formal verification, symbolic execution, and evolutionary algorithms for automated software engineering and scientific discovery.

In summary, SEMREP demonstrates that teaching models to explicitly generate and reason about semantically equivalent code variants is a superior strategy for code transformation, yielding higher correctness, better generalization, and more efficient use of computational resources.