💬 NLP

From Where Words Come: Efficient Regularization of Code Tokenizers Through Source Attribution

This paper introduces Source-Attributed BPE (SA-BPE), a regularization technique that modifies the standard BPE training objective to mitigate token imbalance and overfitting caused by diverse data sources, thereby reducing under-trained tokens in code tokenizers while maintaining inference efficiency and safety.

Original authors: Pavel Chizhov, Egor Bogomolov, Ivan P. Yamshchikov

Published 2026-04-16

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Pavel Chizhov, Egor Bogomolov, Ivan P. Yamshchikov

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 are teaching a robot to read and write computer code. To do this, you first have to teach the robot how to break sentences down into smaller, manageable pieces called tokens. Think of these tokens like LEGO bricks. If you want the robot to build a castle (write code), it needs a box of LEGO bricks that are the right shapes and sizes.

The standard way to make these bricks is called BPE (Byte-Pair Encoding). It's like a machine that looks at a huge pile of code from the internet and says, "Hey, these two letters appear together a lot, let's glue them into one brick!" It keeps gluing the most common pairs together until it has a vocabulary of bricks.

The Problem: The "Junk Drawer" Effect

The paper points out a major flaw in this standard method: The robot ends up with a box full of useless, weird bricks.

Because the training data comes from millions of different code repositories (projects), the standard BPE gets confused by specific, one-off things.

The Analogy: Imagine you are learning English by reading every book ever written. If one author accidentally types "ipsum" instead of "ipsum" in a placeholder text, the standard BPE might decide, "Oh, 'ipsum' is a real word!" and make a special brick for it.
The Result: The robot now has a brick for "ipsum" (with a typo), a brick for a variable name that only existed in one specific project 5 years ago, or a brick for a random string of numbers that never appears again.
Why it matters: These "junk bricks" take up space in the robot's memory. Worse, because the robot never actually uses them to build anything real, it doesn't know what they mean. When the robot sees them later, it might get confused, hallucinate (make things up), or even be tricked by hackers (jailbreaks) who use these weird bricks to bypass safety rules.

The Solution: Source-Attributed BPE (SA-BPE)

The authors propose a new method called SA-BPE. Instead of just asking, "How often does this pair of letters appear?" they ask, "Where does this pair appear?"

They introduce two new rules for making the LEGO bricks:

The "Popularity Contest" (Repository Count): A brick is only allowed if it appears in many different code projects (repositories). If a weird variable name only shows up in one single project, SA-BPE says, "Nope, that's too specific. Don't make a brick for that."
The "Language Passport" (Language Count): A brick is only allowed if it appears in many different programming languages. If a pattern only exists in Python but never in Java or C++, it might be too specific to Python's quirks and shouldn't get a special brick.

The Creative Metaphor:
Think of the standard BPE as a chef who makes a soup by throwing in every ingredient they find in the kitchen, regardless of whether it's a carrot or a rusty spoon. The soup tastes weird and has hard chunks.

SA-BPE is like a strict head chef who says: "We only add ingredients that are found in many different recipes. If a spice only appears in one single, weird recipe from 1995, we leave it out." The result is a cleaner, more versatile soup that tastes good no matter who is eating it.

What Happens When They Use SA-BPE?

The researchers tested this on several coding models (like StarCoder2 and CodeGemma) and found:

Fewer Junk Bricks: The number of useless, "under-trained" tokens dropped dramatically (in some cases, down to just 44 bad tokens out of hundreds of thousands).
Better Compression: The robot can still read the code efficiently. It doesn't need more bricks to do the job; in fact, it uses fewer, better-quality bricks.
No Speed Loss: The best part? The robot reads the code just as fast as before. The new rules only changed how the bricks were made in the factory, not how the robot uses them in the kitchen.

The Big Takeaway

This paper is about cleaning up the "dictionary" of AI models. By making sure the AI only learns words and symbols that are truly common and useful across the entire world of coding (rather than just specific to one project), we get models that are:

Smarter (less likely to hallucinate).
Safer (harder to trick).
More Efficient (less wasted memory).

It's a simple but powerful idea: Don't let the AI memorize the typos and inside jokes of a single project; teach it the universal language of code.

1. Problem Statement

The paper addresses the issue of under-trained tokens in Byte-Pair Encoding (BPE) tokenizers, particularly within the domain of code generation.

The Phenomenon: Standard BPE tokenizers often produce tokens that appear in the vocabulary but are rarely or never seen during model training. These tokens lack adequate embedding representations, leading to "cluttered" vocabularies.
Consequences: Under-trained tokens can cause model hallucinations, serve as vectors for token-level jailbreak attacks, and reduce inference efficiency.
Root Cause in Code: The paper identifies that code tokenizers are prone to overfitting on specific data sources. Common issues include:
- Repository-specific artifacts: Variable names, method names, or placeholder strings (e.g., typos in "Lorem ipsum") unique to a single repository or project.
- Imbalanced diversity: Training data often lacks diversity in repository and language sources, causing the BPE algorithm to prioritize frequent but narrow patterns (e.g., specific CamelCase or snake_case combinations from one large repo) over generalizable patterns.
- Intermediate tokens: Even intermediate tokens used to build larger words can become under-trained if they are part of overfitted, repository-specific merges.

2. Methodology: Source-Attributed BPE (SA-BPE)

The authors propose SA-BPE, a set of modifications to the standard BPE training algorithm that introduces regularization based on data source attribution. Instead of just counting token frequency, SA-BPE tracks the diversity of the sources where token pairs appear.

Core Mechanisms

The method tracks two metrics for every token pair during training:

$R$ (Repositories): The number of distinct code repositories where the pair appears.
$L$ (Languages): The number of distinct programming languages where the pair appears.

The authors introduce two types of modifications to the BPE objective:

A. Merge Skip Criterion (Hard Constraint)
The algorithm skips merging a pair if it does not meet minimum diversity thresholds:

A merge is only executed if the pair appears in at least $R_t$ repositories and $L_t$ languages.
Intuition: If a pair exists in very few sources, it is likely a repository-specific artifact (overfitting) and should be pruned early. This prevents the vocabulary from growing with niche tokens.

B. Priority Criterion (Soft Constraint)
The authors modify the scoring function used to select the next merge. Instead of maximizing raw frequency ( $F$ ), they introduce weighting factors to de-prioritize low-diversity pairs:

$F \cdot L$ : Prioritizes pairs appearing across many languages.
$F \cdot \log(R + 1)$ : Prioritizes pairs appearing in many repositories.
$F \cdot \log(R + 1) \cdot L$ : Combines both factors.
Logic: This approach is conceptually opposite to TF-IDF. While TF-IDF up-weights terms unique to a document, SA-BPE down-weights terms unique to a single repository/language to promote generalization.

Implementation Details

Training: The modifications are applied only during the vocabulary construction phase.
Inference: The inference procedure remains identical to standard BPE. The model uses the resulting vocabulary and merge list without any additional computational overhead during generation.

3. Key Contributions

Analysis of Under-trained Tokens: The paper provides a systematic analysis of under-trained tokens in major code models (CodeGemma, Qwen2.5 Coder, StarCoder2), categorizing them into types like repository-specific variable names, punctuation errors, and intermediate parts of overfitted tokens.
SA-BPE Algorithm: The introduction of source attribution ( $R$ and $L$ ) as a regularization mechanism for BPE training.
Empirical Validation: Training of small Llama-based models (100M parameters) to demonstrate that SA-BPE reduces under-trained tokens without sacrificing compression or requiring changes to the inference pipeline.
Open Source: Release of the sa-bpe implementation.

4. Experimental Results

The authors evaluated SA-BPE on Java, Python, and multilingual datasets (18 languages) using metrics like Compression Rate (CR), Token Coverage, Mean Token Length (MTL), and the count of verified under-trained tokens.

Reduction in Under-trained Tokens:
- In multilingual experiments, the best SA-BPE variant ( $F \cdot \log R \cdot L$ ) reduced the number of verified under-trained tokens from 640 (in standard BPE) to 44.
- This represents a massive reduction in "junk" tokens that the model cannot effectively utilize.
Compression and Coverage:
- SA-BPE tokenizers achieved better or comparable compression rates to standard BPE.
- They showed higher token coverage on unseen evaluation data, indicating better generalization to new repositories.
- The Mean Token Length (MTL) decreased, suggesting the vocabulary is less bloated with long, specific variable names and more focused on reusable components.
Pruning Analysis:
- When pruning vocabularies, SA-BPE tokenizers showed that removing tokens immediately impacted compression (low redundancy), whereas standard BPE could have thousands of tokens removed with no effect on compression (high redundancy of useless tokens).
Comparison to Alternatives:
- SA-BPE outperformed other BPE variants like BoundlessBPE and PickyBPE in terms of compression and under-trained token reduction.
- It maintained the same inference speed as standard BPE, unlike some alternatives that require complex inference procedures.

5. Significance and Impact

Production Readiness: Since SA-BPE only modifies the training phase and leaves the inference procedure unchanged, it is a drop-in replacement for existing BPE pipelines. It does not increase inference latency or memory usage during generation.
Safety and Robustness: By eliminating under-trained tokens, the method reduces the risk of hallucinations and potential security vulnerabilities (jailbreaks) associated with tokens the model has never properly learned.
Efficiency: It creates a more efficient vocabulary where every token has a higher probability of being useful, leading to better model performance with the same vocabulary size.
Generalizability: While focused on code, the authors note the approach can be adapted for natural language by using "documents" or "books" as the source attribution unit, though the specific "language" metric is unique to the code domain.

In conclusion, the paper demonstrates that source diversity is a critical, often overlooked factor in tokenizer training. By explicitly regularizing against overfitting to single repositories or languages, SA-BPE produces cleaner, more robust tokenizers that improve the efficiency and safety of code generation models.