T ri E mbed: Bridge the Gap between Text and Token Indices with Embedding Reparameterization (2025)

Abstract: The current paradigm of language modeling is a two-stage pipeline that first transforms raw text to token indices, where the distribution is then estimated. It inherently discards linguistic relations between tokens during tokenization, creating a fundamental gap. To address this, we propose TriEmbed , a reparameterization method for embeddings that incorporates the morphological relationships inherent in subword tokenizer algorithms. Specifically, by organizing the vocabulary into a Trie structure, we can encode these relations and reparametrize the embeddings, facilitating the recovery of other linguistic relationships during training. Empirical results across various settings demonstrate that TriEmbed outperforms conventional embeddings from the perspective of scaling, while offering more linguistically informative token embeddings.

## introduction in recent years, language modeling has undergone rapid advancements, unlocking unprecedented potential through large-scale scaling. the current paradigm of language modeling is a pipeline consisting of two decoupled components: a tokenizer and a language model. the tokenizer converts raw text, defined in the text space, into a sequence of token indices in the token index space, where the language model can operate in a differentiable manner. however, this transformation also introduces a fundamental gap: tokenization discards the linguistic relations between tokens (e.g., morphological variations such as "run" and "running" or synonymous expressions like "big" and "large"), reducing them to mere numerical representations (i.e. token indices) with only identity relations preserved. common practice in llms disregards the gap, relying on end-to-end training with the next token prediction objective to implicitly recover these lost relationships. however, this process is uncontrollable, possibly leading to problems including representation degeneration @xcite , frequency bias @xcite and so on. some previous work attempts to include semantic relations by pretraining embedding modules @xcite ). yet, structural discrepancies between different representation spaces commonly induce domain shifts in continual training, resulting in worse performance than random initializations @xcite . to bridge the gap, we propose a novel reparameterization method for embedding modules, triembed, which incorporates the morphological relations inherent in tokenizer algorithms. subword tokenizer builds on word formation morphology. and its training process reflects morphological relations within vocabulary. we encode these relations in ancestor-descendant connections within an implicit trie. following the hierarchical structure of the trie, we derive each embedding from all its ancestors. our method retains the benefits of random initialization while seamlessly integrating inductive bias about the morphological relations among token indices. we conduct extensive experiments to validate the effectiveness of triembed across different scales of training corpora and model sizes. notably, when pretraining pythia on fineweb-edu, triembed achieves performance on par with conventional embedding modules while requiring about 70% of the model parameters or 90% of the pretraining corpus. furthermore, a detailed analysis reveals that our approach significantly improves the morphology of representation space, offering more meaningful token embeddings. ## background language modeling aims to estimate the probability p(s) of any given text string @xmath0. current paradigm follows a two-component pipeline consisting of a tokenizer and a language model. a tokenizer segments a text string s into tokens, then maps them to a sequence of token indices @xmath1, ...x n ] ∈ z * . it provides an intermediate token index space, in which the language model can estimate the probability distribution p θ (x i |@xmath2. tokenizer algorithm focuses on constructing a suitable vocabulary. llms predominantly adopt subword tokenizer such as bpe @xcite , which hypothesizes that complex words should be split into multiple subwords. for training, they begin with a base vocabulary (e.g. all unicode characters), and iteratively merges the most frequent token pairs in training corpus as a new token to expand the vocabulary (see appendix a for the pseudo-code). therefore, the resulting vocabulary implicitly encodes a trie foot_0 , as shown in figure 1(a) . although the pipeline simplifies the training process for next token prediction, it also introduces a fundamental gap when transforming from the original text space to the token index space. to illustrate this issue, consider the word spelling challenge. a sample @xmath3, id(→ ), id(_m), id(_o), id(_d), id(_e), id(_l)]. the morphological relations between different tokens (e.g. "model" and "_m") are lost, leading to frequent failure of llms in word spelling tasks @xcite . this issue extends beyond morphology and is inherent to the two-component pipeline. a wide range of linguistic relationships, including semantic, etymological, and derivational connections, are similarly discarded after tokenization, since lan-guage models can only observe token indices with identity relations left. common practice in llms tends to disregard the gap, relying on end-to-end training with the next token prediction objective to learn these linguistic information implicitly. however, this process is computational expensive. moreover, the stochastic nature of gradient-based algorithm makes the training process uncontrollable, leading to severe problems such as representation degeneration @xcite and frequency bias @xcite @xcite . another line of work @xcite tries to initialize embeddings with pretrained weights to partially recover the missing relations among tokens in vocabulary. however, these pretrained embeddings often exhibit structural discrepancy with the final representation space of the language model. for instance, they tend to have significantly larger variance magnitudes @xcite , which makes the subsequent training difficult. more related work is presented in appendix b. ## method as discussed in the previous section, the transformation from text space to token index space discards linguistic relationships between tokens, introducing a fundamental gap. to address this, we aim to leverage the inductive bias inherent in the tokenizer to recover the lost information during tokenization. the subword tokenizer draws inspiration from word formation morphology @xcite @xcite , a field that studies the structure of words in terms of morphemes. and its iterative merging process mirrors the word formation process where words are constructed by sequentially adding affixes to stems. as such, an inherent morphological relation exists among the resulting vocabulary. by reorganizing the vocabulary into a trie structure, this morphological information is captured within the ancestor-descendant relationships of the tree. we hypothesize that, the ancestor-descendant relationship can serve as a foundation for the intricate linguistic relations among tokens, facilitating the restoration of other relations during training. building on the hypothesis, we propose triembed, a novel reparameterization approach for the embedding module, which incorporates the inductive bias about the trie structure among tokens. given a token, we derive its embedding from those of its ancestors. specifically, we first assign each token with a unique atomic vector a(•), then define the embedding of a token as the cumulative sum of the atomic vectors of all its ancestors and its own, i.e. e(v) = @xmath4. this summation process follows a traversal from the root of the trie to the current node, reflecting the word formation process in which affixes are progressively added to a stem. the method is illustrated in figure 1(b) . a key advantage of our approach is that it does not introduce additional parameters, maintaining parameter efficiency. additional computational costs are also minimal due to simplicity of summation. despite its simplicity, triembed still brings significant improvements, which we will empirically demonstrate in the following experiments. ## conclusions in this work, we review the gap between text space and token index space inherent in the current twocomponent language modeling pipeline. we introduce triembed, a reparameterization method for the embedding module that incorporates morphological information from the tokenizer algorithm. extensive experiments demonstrate the effectiveness of triembed particularly in terms of scalability, suggesting a promising direction for future large-scale pretraining of language models.