LLMs as Code Architects: Meta’s New Approach to Precise Code Transformations

Tools designed for rewriting, refactoring, and optimizing code should prioritize both speed and accuracy. Large language models (LLMs), however, often lack these critical attributes. Despite these limitations, there remains significant potential for leveraging LLMs to enhance code quality.

In a new paper Don’t Transform the Code, Code the Transforms: Towards Precise Code Rewriting using LLMs, a Meta research team proposes a novel chain-of-thought strategy to efficiently generate code transformations using LLMs. Their approach enables LLMs to derive transformations based on a small set of input/output examples.

A code transformation refers to a function that modifies existing code to achieve a desired form, a method applicable to various tasks—from compiler optimizations to legacy code refactoring. Traditional rule-based transformations, while effective, are often complex to design and implement. LLMs present a promising alternative, but their logic can be opaque, lacks guarantees of correctness, and is difficult to debug when errors occur.

Instead of using LLMs to directly rewrite code, the researchers propose a system where LLMs serve as generators for code transformations. Their method hinges on a step-by-step approach, using introspective iterations to refine the model’s output. This process allows the LLM to diagnose its mistakes and hypothesize potential solutions before making adjustments.

The process involves several key steps:

1. Initial Example Presentation: The model is given three input/output examples and tasked with describing the underlying transformation in clear natural language.
2. Iterative Refinement: The model refines its description, if necessary, over 10 iterations. The initial attempts often fail to capture the full complexity of the transformation, especially for edge cases.
3. Generating the Transformation: Once the description is deemed sufficient, the model uses it, along with the original examples, to produce a code transformation implementation.
4. Testing and Debugging: The generated transformation is then tested in a controlled environment with 10 input/output examples (including the initial ones). If errors occur or the output is incorrect, counterexamples or error messages are provided to the model.
5. Failure Analysis: When failures occur, the model performs an introspection, analyzing the error and attempting to explain the issue.
6. Iterative Fixing: The model is prompted with the incorrect code, the failure case, and its own analysis, iterating until a working transformation is produced (up to 50 iterations). Afterward, the quality of the transformation is evaluated on unseen examples.

The research team compared their Code the Transforms (CTT) method with a Transform the Code (TTC) approach, where the LLM directly rewrites code. Their findings show that CTT consistently outperforms TTC in precision, achieving an accuracy rate of 0.95 compared to TTC’s 0.60. Furthermore, errors in CTT are generally easier to diagnose and resolve. On average, the synthesized transformations consisted of 34.4 lines of code, whereas TTC required extensive review to catch mistakes.

In conclusion, the paper demonstrates that having LLMs generate transformations rather than directly rewriting code not only increases precision but also reduces the cost of reviewing and testing code, offering a more efficient alternative for code optimization tasks.

The paper Don’t Transform the Code, Code the Transforms: Towards Precise Code Rewriting using LLMs is on arXiv.

---

Author: Hecate He | Editor: Chain Zhang