Stable Code Technical Report (2404.01226v1)

Abstract: We introduce Stable Code, the first in our new-generation of code LLMs series, which serves as a general-purpose base code LLM targeting code completion, reasoning, math, and other software engineering-based tasks. Additionally, we introduce an instruction variant named Stable Code Instruct that allows conversing with the model in a natural chat interface for performing question-answering and instruction-based tasks. In this technical report, we detail the data and training procedure leading to both models. Their weights are available via Hugging Face for anyone to download and use at https://huggingface.co/stabilityai/stable-code-3b and https://huggingface.co/stabilityai/stable-code-instruct-3b. This report contains thorough evaluations of the models, including multilingual programming benchmarks, and the MT benchmark focusing on multi-turn dialogues. At the time of its release, Stable Code is the state-of-the-art open model under 3B parameters and even performs comparably to larger models of sizes 7 billion and 15 billion parameters on the popular Multi-PL benchmark. Stable Code Instruct also exhibits state-of-the-art performance on the MT-Bench coding tasks and on Multi-PL completion compared to other instruction tuned models. Given its appealing small size, we also provide throughput measurements on a number of edge devices. In addition, we open source several quantized checkpoints and provide their performance metrics compared to the original model.

• The paper introduces Stable Code, a new language model variant optimized for code tasks with a novel multi-stage training approach.
• It employs a decoder-only transformer enhanced by techniques like Rotary Position Embeddings and the FIM objective to improve code completion and fill-in-the-middle tasks.
• The model’s lightweight design enables efficient use on consumer-grade hardware, reducing latency and supporting real-time code interactions.

Stable Code Technical Report

Introduction

The "Stable Code Technical Report" introduces Stable Code and its variant Stable Code Instruct, which are part of a new generation of LLMs optimized for code-related tasks. These models aim to enhance capabilities in code completion, reasoning, and related software engineering tasks. A crucial aspect of this work is its focus on creating lightweight models capable of running efficiently on edge devices.

Model Architecture and Training

Stable Code is built on Stable LM 3B, which employs a decoder-only transformer architecture. The model leverages architectural modifications, including Rotary Position Embeddings and LayerNorm, ensuring performance optimization for code-related tasks. A significant innovation in the training process is the adoption of a multi-stage approach, which leverages a pre-training dataset derived from a diverse array of programming languages and technical documents (Figure 1).

Figure 1: Staged approach to training Stable Code 3B and Stable Code Instruct 3B.

The training employs advanced techniques such as AdamW optimizer and FIM objective. This incorporation of FIM helps the model predict code more contextually, reflecting the non-linear nature of code execution and structure. The staged approach is complemented by fine-tuning stages like SFT and DPO to further refine the model's instruction-based interactions.

Performance Evaluation

Stable Code 3B and its instruct variant have been extensively evaluated across various benchmarks:

1. Code Completion: Stable Code demonstrates competitive performance on the Multi-PL benchmark, rivaling larger models such as Code Llama and StarCoder 15B, despite its smaller size.
2. Fill in the Middle (FIM) Task: The models excel in FIM tasks—showcasing improved prediction capabilities on benchmarks like StarCoder-FIM, illustrating enhanced understanding and completion abilities in non-linear code contexts (Figure 2).

   Figure 2: Stable Code 3B Loss and Learning Rate Curves.

3. Instruction Tuning: Stable Code Instruct outperforms similar scale models in instruction-based tasks, particularly evident in complex multi-turn interactions, such as those found in MT-Bench.

   Figure 3: Code Performance Comparison of Stable Code 3B Scratch and Stable LM 3B Initializations.

Practical Implications

Stable Code's development underscores the potential for creating efficient, high-performance models for software engineering applications. With its ability to function effectively on consumer-grade hardware, it facilitates a reduction in latency and dependencies often associated with cloud-based solutions. This architectural efficiency is critical for applications requiring real-time code completions and interactions.

Conclusion

The launch of Stable Code and its instruct version marks a strategic advancement in AI-driven code modeling. Their robust performance, especially in handling multilingual and multi-turn code tasks, presents promising avenues for extending LLM applications in software development environments. The open-source release signifies an opportunity for the broader research community to innovate upon and tailor these models for diverse and specialized code solutions. The exploration of quantized weights represents a step forward in optimizing inference speed and resource utilization, making these models highly adaptable for scalable deployment scenarios.