Enhancing Cross-Language Code Translation via Task-Specific Embedding Alignment in Retrieval-Augmented Generation

TL;DR

The paper tackles the challenge of cross-language code translation from Fortran to C++ by introducing task-specific embedding alignment within a Retrieval-Augmented Generation (RAG) framework. It trains a Fortran embedding model with a Soft Information Noise-Contrastive Estimation (S-InfoNCE) loss guided by CodeBLEU similarities between generated translations, aligning retrieval to improve downstream translation quality without fine-tuning the LLM. Empirical results on HPC Fortran2CPP and Numerical Recipes show substantial CodeBLEU gains (e.g., 0.64→0.73 and 0.52→0.60 respectively), demonstrating the approach's effectiveness and practical efficiency across multiple model sizes. The method leverages a 25k-Fortran snippet dataset from Stack-V2 and translations produced by LLaMA 3.1-8B, highlighting the potential to boost code translation performance in resource-constrained settings and extending to other language pairs and task-specific metrics.

Abstract

We introduce a novel method to enhance cross-language code translation from Fortran to C++ by integrating task-specific embedding alignment into a Retrieval-Augmented Generation (RAG) framework. Unlike conventional retrieval approaches that utilize generic embeddings agnostic to the downstream task, our strategy aligns the retrieval model directly with the objective of maximizing translation quality, as quantified by the CodeBLEU metric. This alignment ensures that the embeddings are semantically and syntactically meaningful for the specific code translation task. Our methodology involves constructing a dataset of 25,000 Fortran code snippets sourced from Stack-V2 dataset and generating their corresponding C++ translations using the LLaMA 3.1-8B language model. We compute pairwise CodeBLEU scores between the generated translations and ground truth examples to capture fine-grained similarities. These scores serve as supervision signals in a contrastive learning framework, where we optimize the embedding model to retrieve Fortran-C++ pairs that are most beneficial for improving the language model's translation performance. By integrating these CodeBLEU-optimized embeddings into the RAG framework, our approach significantly enhances both retrieval accuracy and code generation quality over methods employing generic embeddings. On the HPC Fortran2C++ dataset, our method elevates the average CodeBLEU score from 0.64 to 0.73, achieving a 14% relative improvement. On the Numerical Recipes dataset, we observe an increase from 0.52 to 0.60, marking a 15% relative improvement. Importantly, these gains are realized without any fine-tuning of the language model, underscoring the efficiency and practicality of our approach.

Figures (3)

• Figure 1: Overview of the proposed pipeline. i) The LLM generates pairwise code translations, which are evaluated using the CodeBLEU metric. ii) The resulting similarity scores are used to guide contrastive learning for semantic alignment of the embedding model.
• Figure 2: Scatter plots comparing the unaligned and aligned One-shot CodeBLEU scores across different shot counts (1-shot, 2-shot, 3-shot) for two models (llama3.1 70b and llama3.1 8b) and two datasets (Numerical Recipe and HPC Fortran2C++ Dataset). Each point represents a shot count, and the red dashed line represents the reference where the unaligned and aligned scores are equal. The text box in each subplot displays the average CodeBLEU performance and standard deviation for aligned vs. unaligned RAG translation across the few-shot configurations.
• Figure 3: Box plots illustrating the distribution of CodeBLEU scores across various shot counts (1-shot, 2-shot, 3-shot) for both unaligned and aligned models. The results are presented for two models (llama3.1 70b and llama3.1 8b) across two datasets (Numerical Recipe and HPC Fortran2C++ Dataset)

Theorems & Definitions (2)

• Lemma 1
• proof