Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Adapting Language Models via Token Translation

Zhili Feng, Tanya Marwah, Nicolo Fusi, David Alvarez-Melis, Lester Mackey

TL;DR

S2T2 improves both the perplexity and the compression of out-of-domain protein sequences and finds that token translations learned for smaller, less expensive models can be directly transferred to larger, more powerful models to reap the benefits of S2T2 at lower cost.

Abstract

Modern large language models use a fixed tokenizer to effectively compress text drawn from a source domain. However, applying the same tokenizer to a new target domain often leads to inferior compression, more costly inference, and reduced semantic alignment. To address this deficiency, we introduce Sparse Sinkhorn Token Translation (S2T2). S2T2 trains a tailored tokenizer for the target domain and learns to translate between target and source tokens, enabling more effective reuse of the pre-trained next-source-token predictor. In our experiments with finetuned English language models, S2T2 improves both the perplexity and the compression of out-of-domain protein sequences, outperforming direct finetuning with either the source or target tokenizer. In addition, we find that token translations learned for smaller, less expensive models can be directly transferred to larger, more powerful models to reap the benefits of S2T2 at lower cost.

Adapting Language Models via Token Translation

TL;DR

S2T2 improves both the perplexity and the compression of out-of-domain protein sequences and finds that token translations learned for smaller, less expensive models can be directly transferred to larger, more powerful models to reap the benefits of S2T2 at lower cost.

Abstract

Modern large language models use a fixed tokenizer to effectively compress text drawn from a source domain. However, applying the same tokenizer to a new target domain often leads to inferior compression, more costly inference, and reduced semantic alignment. To address this deficiency, we introduce Sparse Sinkhorn Token Translation (S2T2). S2T2 trains a tailored tokenizer for the target domain and learns to translate between target and source tokens, enabling more effective reuse of the pre-trained next-source-token predictor. In our experiments with finetuned English language models, S2T2 improves both the perplexity and the compression of out-of-domain protein sequences, outperforming direct finetuning with either the source or target tokenizer. In addition, we find that token translations learned for smaller, less expensive models can be directly transferred to larger, more powerful models to reap the benefits of S2T2 at lower cost.

Paper Structure

This paper contains 7 sections, 1 equation, 2 figures, 1 table, 2 algorithms.

Table of Contents

  1. Introduction
  2. Translating Tokens with Sparse Sinkhorn
  3. Finding ${\mathbf{P}}$ via Sparse Sinkhorn
  4. Experiment
  5. Training details.
  6. Conclusion
  7. Acknowledgement

Figures (2)

  • Figure 1: Overview of S2T2. Left: S2T2 injects a weight-tied sparse optimal transport (OT) layer in both the token embedding and language model head. The input tokens will be encoded based on a sparse convex combination of the original token embeddings and decoded by a sparse combination of the original language model head. Right: The sparse OT matrix is obtained by iteratively projecting a dense cost matrix along its rows and columns. The dense cost matrix is updated by backpropogation.
  • Figure 2: Evaluation loss after initializing OLMo-7B with token translator ${\mathbf{P}}$ learned from OLMo-1B. Along the x-axis, $\text{S2T2-}\alpha$ represent S2T2 with the $\alpha$-entropy regularizer that controls the sparsity of ${\mathbf{P}}$. New Tok. is OLMo-7B with the new tokenizer and truncated ${\mathbf{E}}, {\mathbf{L}}$; Orig Tok. is OLMo-7B with the original tokenizer. The red dashed line is the loss when you randomly guess the next token.