🤖 AI

MoSE: Hierarchical Self-Distillation Enhances Early Layer Embeddings

The paper introduces MoSE, a 1-billion-parameter multi-exit encoder that utilizes hierarchical self-distillation and repository-level contextual loss to enhance early-layer embeddings, enabling flexible and cost-effective deployment for code retrieval and classification tasks.

Original authors: Andrea Gurioli, Federico Pennino, João Monteiro, Maurizio Gabbrielli

Published 2026-01-27

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Andrea Gurioli, Federico Pennino, João Monteiro, Maurizio Gabbrielli

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a giant, incredibly smart library assistant (a Large Language Model) who knows everything about computer code. This assistant is brilliant, but it's also huge, slow, and requires a massive amount of electricity to run. If you want to use it on a regular laptop or a phone, it's just too heavy.

Usually, to make this assistant smaller, researchers would try to "distill" it—like taking a big pot of soup and trying to boil it down to a tiny cup while keeping the flavor. This is expensive and often loses some of the taste (accuracy).

MoSE is a new way of building this assistant that solves this problem without needing a separate "teacher" to guide it. Here is how it works, using some simple analogies:

1. The "Multi-Exit" Elevator

Think of a standard AI model as a 36-story building. In a normal building, you have to go all the way to the top floor (Layer 36) to get the "best answer." If you stop on the 4th floor, the answer is usually weak or wrong.

MoSE is different. It's like an elevator with five special exits (at floors 4, 9, 18, 27, and 36).

The Problem: Usually, the lower floors are messy and don't know much.
The MoSE Solution: The architects (the researchers) taught the building so that every floor knows how to give a good answer, not just the top one.
How? They used a trick called Self-Distillation. Imagine the smart people on the top floor (Layer 36) are constantly whispering advice to the people on the lower floors (Layers 4, 9, etc.). By the time the lower floors finish their work, they are almost as smart as the top floor.

2. The "Speed vs. Accuracy" Switch

Because every floor is now smart, you get to choose your speed:

Need it fast? You stop the elevator at Floor 4. It uses very little energy and gives an answer almost instantly. It's 90% faster than going to the top, and the answer is still very good (only about 6% less accurate).
Need it perfect? You go all the way to Floor 36.
The Magic: You don't need to train five different models. You have one model that can act like five different sizes, depending on how much time and battery you have.

3. Learning from the Whole Neighborhood (Context)

Most code models learn by looking at one file at a time, like reading a single page of a book in isolation.

MoSE's Approach: It looks at the whole repository (the whole neighborhood of files).
The Analogy: Instead of just reading a single sentence, MoSE reads a whole chapter to understand the context. It asks, "Do these two code snippets belong to the same project?" This helps it understand the "vibe" of the code better, even if the code is written in different languages (like translating Python to Rust).
The Result: They created a new dataset called SynthCoNL where they took code snippets and translated them into different languages to teach the model that "this Python code means the same thing as that Rust code."

4. Why This Matters

Efficiency: You can run a powerful code-search tool on a small device without needing a supercomputer.
Flexibility: If you are building a simple app, you use the "early exit" (small, fast). If you are doing complex research, you use the "deep exit" (big, thorough).
No Extra Cost: Unlike old methods that required training a huge model first and then shrinking it, MoSE learns to be efficient while it is being trained.

In short: MoSE is a smart, modular code assistant that lets you choose how much "brainpower" you want to use. It learns from its own "older, smarter self" (the deeper layers) to make sure even its "younger, faster self" (the early layers) is ready to do the job.

Technical Summary: MoSE – Hierarchical Self-Distillation Enhances Early Layer Embeddings

1. Problem Statement

Deploying Large Language Models (LLMs) in production often necessitates navigating a trade-off between accuracy and performance to meet strict latency and resource constraints. Traditional model distillation reduces model size but incurs significant costs by requiring the training of separate "student" models to emulate "teacher" models. Furthermore, existing approaches often assume that the final layer of a transformer produces the most useful representations, whereas recent insights suggest that intermediate layers may hold the most semantically rich information for specific downstream tasks. There is a need for a unified architecture that allows flexible deployment (selecting different model portions) without sacrificing performance or requiring separate distillation pipelines.

2. Methodology

The authors introduce MODULARSTARENCODER (MOSE), a 1-billion-parameter multi-exit encoder built upon the StarCoder-2 architecture. The core innovation is a Self-Distillation mechanism that enhances lower-layer representations by training multiple exit points simultaneously within a single model stack.

Architecture and Training Strategy

Multi-Exit Design: The model features 36 hidden layers with exit heads at layers 4, 9, 18, 27, and 36. During inference, users can select an exit point based on computational constraints (e.g., stopping at layer 4 for speed or layer 36 for maximum accuracy).
Self-Distillation Loss: The training objective aggregates losses from all selected exit layers. Higher layers guide earlier ones, encouraging lower layers to learn robust representations. The total loss is computed as a weighted sum: $L = \sum_{i \in \iota} L_i \cdot \alpha_i$ , where $\alpha_i$ increases with layer depth to prioritize deeper layers.
Pre-training Objectives:
- Masked Language Modeling (MLM): Standard token prediction.
- In-Context Classification (ICC): Replaces the traditional Next Sentence Prediction (NSP) loss. ICC classifies whether concatenated code snippets originate from the same repository. This increases input density (reducing padding), expands the context window utilization, and fosters better chunk-level comprehension essential for semantic retrieval.
Contextual Enhancements: The model utilizes Grouped Query Attention (GQA), Rotary Positional Encoding (RoPE), and FlashAttention V2. The causal masking of StarCoder-2 was removed to enable bidirectional attention, and sliding window attention was replaced with full attention to avoid receptive field limitations.

Data and Fine-tuning

SYNTHCoNL Dataset: A new dataset of 1.07 million natural language-code-code triplets. It was constructed by translating code snippets from the CodeSearchNet dataset into multiple languages (Python, Java, Go, PHP, etc.) using Qwen2.5-Coder-7B-Instruct. This creates cross-language code-to-code pairs tied to the same natural language description, expanding text-to-code benchmarks.
Fine-tuning: MOSE was fine-tuned on SYNTHCoNL for text-to-code and code-to-code retrieval using a CLIP-style loss with multi-layer projection heads. It was also fine-tuned on BigCloneBench for code clone detection.

3. Key Contributions

Self-Distillation Methodology: A novel approach that trains multiple models within a unified layer stack, eliminating the need for separate teacher/student models and reducing redundancy.
MOSE Architecture: A pre-trained and fine-tuned 1B-parameter encoder with five exit points, allowing users to dynamically select model size based on memory and computational constraints while maintaining competitive performance.
SYNTHCoNL Dataset: A new benchmark resource containing 1,071,367 triplets generated via code translation, facilitating cross-language code-to-code retrieval research.

4. Results

Evaluations were conducted on the CodeXGLUE benchmark suite (CodeSearchNet, Code Translation, POJ104) and BigCloneBench.

Text-to-Code Retrieval (CodeSearchNet): MOSE achieved an average Mean Reciprocal Rank (MRR) of 81.0, outperforming all open-source models, including CodeT5+ (770M), UniXcoder, and ModernBERT-Large. It narrowed the performance gap with the closed-source OpenAI Text-Embedding-3-Large.
Code-to-Code Retrieval:
- On the POJ104 dataset (intra-lingual semantic search), MOSE achieved a mean Average Precision (mAP) of 75.9 in a zero-shot setting, the highest among open-source models.
- On the Code Translation (CT) dataset, MOSE performed comparably to OpenAI's embedding model and outperformed other open models.
Code Clone Detection (BigCloneBench): MOSE achieved an F1 score of 94.2 at Layer 18, matching the performance of deeper layers and competing with CodeT5+ (95.1) and UniXcoder (95.2).
Efficiency-Accuracy Trade-off: The paper demonstrates that reducing the model from 36 layers to 4 layers results in a 90% reduction in floating-point operations (GFLOPs) with only a 6.4% absolute drop in MRR for text-to-code retrieval. Notably, for code-to-code tasks, the best performance was often observed in the initial layers (e.g., Layer 4), challenging the assumption that deeper layers are always superior.

5. Significance and Claims

The paper claims that MOSE offers a principled approach to trading inference cost for accuracy. By leveraging Self-Distillation, the model enhances intermediate representations, making early layers viable for deployment without the overhead of traditional distillation.

Deployment Flexibility: The architecture provides a natural model selection mechanism where users can choose the exit point that best fits their latency and energy constraints.
Representation Divergence: The results highlight that the optimal layer for semantic alignment depends on the input modality (text-to-code vs. code-to-code), with shallow layers often performing better for code-to-code tasks.
Efficiency: The approach supports sustainable deployment by enabling smaller, energy-efficient model configurations without significant performance degradation.

The authors acknowledge limitations, including computational constraints that limited extensive hyperparameter tuning and the unclear impact of using synthetic code generated via translation for fine-tuning. Future work is proposed to further explore the interactions between task types and layer depth.