On the Effect of Token Merging on Pre-trained Models for Code

2025-07-18

AI Summary

This work addresses the computational overhead induced by subword representation redundancy in code pre-trained models. We propose two semantic-driven hidden-layer representation merging strategies: static averaging and dynamic fusion via learnable weights. To our knowledge, this is the first systematic investigation into merging hidden representations of subwords belonging to the same semantic unit—e.g., constituent subwords of a single identifier—while preserving identifier-level semantic integrity. Our approach integrates seamlessly into mainstream models including CodeBERT, UniXCoder, and CodeT5+, and is empirically validated on vulnerability detection, code classification, and code translation tasks. Experiments demonstrate a 1–19% reduction in inference computation, a +2.47 improvement in CodeBLEU for code translation, and only a marginal −1.82 drop in F1-score for vulnerability detection—achieving a Pareto improvement in both efficiency and performance.

Technology Category

• Natural Language Processing: Code Generation / Program Synthesis from Natural Language
• Machine Learning: Representation Learning
• Knowledge Representation and Reasoning: Computational Complexity of Reasoning

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Representation, reconstruction, and subgraph or motif discovery in Web-related graphs
• Search and Retrieval-Augmented AI: Web query analysis, representation and understanding

Abstract

Tokenization is a fundamental component of language models for code. It involves breaking down the input into units that are later passed to the language model stack to learn high-dimensional representations used in various contexts, from classification to generation. However, the output of these tokenizers is often longer than that traditionally used in compilers and interpreters. This could result in undesirable effects, such as increased computational overhead. In this work, we investigate the effect of merging the hidden representations of subtokens that belong to the same semantic unit, such as subtokens that form a single identifier. We propose two strategies: one based on averaging the representations and another that leverages a learning-based approach. Both methods can be seamlessly integrated with existing language models for code. We conduct experiments using six language models for code: CodeBERT, GraphCodeBERT, UniXCoder, CdoeT5, CodeT5+ (220M), and CodeT5+ (770M), across three software engineering tasks: vulnerability detection, code classification, and code translation. Results show that these strategies can reduce the number of floating-point operations by $1%$ to $19%$. Regarding downstream performance, the most significant degradation was observed in the vulnerability detection task, where the F1 score decreased by $1.82$ points compared to the baseline. In contrast, for code translation, we observed an improvement of $2.47$ points in CodeBLEU. This work contributes to the broader effort of improving language models for code across multiple dimensions, including both computational efficiency and downstream performance.

Problem

• Investigates token merging impact on code language models
• Proposes strategies to reduce computational overhead in tokenization
• Evaluates methods on multiple code tasks and models

Innovation

• Merge subtokens via averaging representations
• Use learning-based approach for token merging
• Integrate methods with existing code models