---
title: Sparse Sinkhorn Token Translation (S2T2)
url: https://www.emergentmind.com/topics/sparse-sinkhorn-token-translation-s2t2
type: topic
---

# Sparse Sinkhorn Token Translation (S2T2)

Sparse Sinkhorn Token Translation (S2T2) is an adaptation technique for pre-trained language models that leverages domain-specific tokenizers and structured, sparse token translation plans. It enables efficient reuse of next-token predictors across domains with differing token statistics, such as from natural language to protein sequences, by learning a transport matrix between token vocabularies, subject to marginal constraints and regularized for sparsity. S2T2 formalizes token translation as an entropic regularized optimal transport problem solved via the Sinkhorn-Knopp algorithm, producing a soft, potentially sparse mapping between target and source tokens to maximize model performance and compression in out-of-domain tasks [2411.00593].

## 1. Problem Statement and Notation

Let $V_s$ be the source vocabulary (e.g., English byte-pair encoding) of size $|V_s| = v$, for which a language model $L_s$ has been pre-trained to predict the next token via cross-entropy. For a new target domain where $V_s$ does not yield efficient compression or semantic alignment—common in non-natural-language domains like proteins—a new tokenizer $T_t$ is trained, producing a target vocabulary $V_t$ of size $|V_t| = u$. The adaptation problem is: reuse $L_s$ (fixed weights) but operate over $V_t$ by translating each target token $t_i \in V_t$ into a sparse distribution over source tokens in $V_s$ (“soft” tokens), feed these through $L_s$ to predict the next source token, and translate the predicted source-token distribution back to $V_t$. This bidirectional translation is achieved with a learned matrix $M \in \mathbb{R}^{u \times v}$ interpreted as a transport plan between $V_t$ and $V_s$ [2411.00593].

## 2. Mathematical Foundations of Sparse Sinkhorn Token Translation

The translation matrix $M$ is directly parameterized in $\mathbb{R}^{u \times v}_{\geq 0}$ and constrained such that:
\[
\sum_{j=1}^{v} M_{ij} = \alpha_i \quad \forall i \in V_t \qquad
\sum_{i=1}^{u} M_{ij} = \beta_j \quad \forall j \in V_s
\]
with frequency marginals $\alpha, \beta$ typically uniform ($\alpha_i = 1/u$, $\beta_j = 1/v$).

The S2T2 objective jointly optimizes tokenizer parameters $\theta_t$ and $M$:
\[
L(M, \theta_t) = \mathbb{E}_{x\sim\text{Data}}\big[-\log P_{L_s}(\operatorname{sink}_\epsilon(M)\odot T_t(x) \to \text{next-token})\big] + \lambda \Omega(M)
\]
where $T_t(x)$ is the one-hot encoding of a target sequence, $\operatorname{sink}_\epsilon(M)$ applies a regularized Sinkhorn projection, $P_{L_s}$ is the fixed LM probability, and $\Omega(M)$ enforces sparsity (e.g., $\ell_1$ or group-L1 penalty).

The entropic regularized optimal transport solution is:
\[
\operatorname{sink}_\epsilon(C) = \arg\min_{P\in U(\alpha, \beta)} \sum_{i,j} P_{ij} C_{ij} + \epsilon \sum_{i,j} P_{ij} (\log P_{ij} - 1)
\]
with $C$ as a cost/score matrix and $P$ produced by Sinkhorn-Knopp iterations:
- Initialize $K = \exp(-C/\epsilon)$ elementwise.
- Iterate: $u^{(t+1)} = \alpha ./ (K v^{(t)})$, $v^{(t+1)} = \beta ./ (K^T u^{(t+1)})$.
- Form $P = \operatorname{diag}(u^{(T)}) K \operatorname{diag}(v^{(T)})$.
After convergence, $P$ satisfies prescribed marginals. Sparsity regularization is implemented by backpropagating through Sinkhorn steps into $C$ and hence $M$.

## 3. Optimization and Implementation Procedure

S2T2 parameters $(\theta_t, C)$ are optimized end-to-end with $T$ unrolled Sinkhorn iterations per forward pass. The workflow is:

1. Sample target-domain data $x$, tokenize with $T_t$ and encode as $T_t(x) \in \{0,1\}^{n\times u}$.
2. Compute Sinkhorn-projected transport $P = \operatorname{sink}_\epsilon(C)$.
3. Feed “soft” source tokens $M^T T_t(x)$ through the fixed $L_s$ to calculate next-token cross-entropy.
4. Add $\lambda \Omega(P)$ to penalize translation density.
5. Backpropagate through Sinkhorn iterations to update $C$ and tokenizer merges.
6. Hyperparameters: regularization weight $\lambda$ (0–1), Sinkhorn temperature $\epsilon = 0.05$–$0.1$, $T = 3$ Sinkhorn iterations, batch size $16$, sequence length $512$, AdamW ($\text{lr}(M)=10^{-3}$, $\text{lr}(\theta_t)=2\times 10^{-5}$), cosine-annealed schedule with 20% warm-up [2411.00593].

## 4. Empirical Performance and Benchmarking

S2T2 is evaluated on adaptation from English to protein sequence modeling (UniRef50). A new BPE tokenizer $T_t$ (vocab size $u = 512$) reduces token length by $\sim$1.82$\times$. Performance is measured via:
- Perplexity (“perp”): LM accuracy conditioned on token translations.
- Bits per byte (BpB): $-\log_2 P(t|\,\text{context})/\lvert\text{original bytes}\rvert$.

Comparison with baselines (unconstrained M, dense Sinkhorn, finetuning with source/target tokenizers) is summarized:

| Method                     | Perplexity ↓ | BpB ↓  |
|----------------------------|--------------|--------|
| Unconstrained M            |   174.20     | 4.09   |
| + continual finetune       |   130.44     | 3.86   |
| Dense Sinkhorn M           |   167.74     | 4.06   |
| + continual finetune       |   136.12     | 3.89   |
| S2T2 (sparse Sinkhorn M)   |   144.03     | 3.94   |
| + continual finetune       | **118.78**   | **3.78**|
| Finetune, orig. tokenizer  |   151.05     | 7.24   |
| Finetune, new tokenizer    |   130.56     | 3.86   |

S2T2 combined with LM finetuning yields the lowest reported perplexity and best compression, outperforming both naive and dense translation strategies [2411.00593].

## 5. Transferability Across Model Scales

A notable property of S2T2 is the portability of the learned translation matrix $M$ across models of different sizes. Matrices trained on a 1B-parameter LLM (OLMo-1B) can be rescaled and directly plugged into larger models (OLMo-7B) without retraining $M$. Experiments show:
- Direct transfer and initialization with fixed $M$ provides immediate perplexity and compression gains relative to both original and target-domain tokenizers.
- Learning $M$ on a small model is approximately 7$\times$ cheaper than training the same on the larger model.
- Figure 2 in [2411.00593] shows that transferred S2T2 improves held-out loss from the outset, regardless of sparsity penalty $\lambda$.

This suggests translation plans learned for token alignment are robust and reusable across model scales.

## 6. Advantages, Limitations, and Prospects

S2T2’s main advantages are:
- Domain-specific tokenizer learning with preserved pre-trained LM capabilities.
- Empirically superior perplexity and data compression compared to direct finetuning with static tokenizers.
- Sparser translations enable efficient inference, as most target tokens map to a small number of source tokens.
- Translation matrices are readily transferable across model scales, significantly reducing computational cost in adaptation [2411.00593].

Noted limitations and open questions include:
- Computational overhead of Sinkhorn iterations and need to carefully tune temperature $\epsilon$ and number of iterations.
- The choice of the sparsity penalty $\Omega(M)$ strongly affects the trade-off between translation matrix density and accuracy.
- Extension to multi-domain or multilingual adaptation would require block-structured, potentially hierarchical $M$.
- Adapting S2T2 to modalities beyond text (e.g., vision, code) would involve defining suitable “tokenizers” aligned with the modality structure.

## 7. Connections to Sparse Sinkhorn Attention and Broader Context

S2T2 incorporates entropic-regularized transport and Sinkhorn projections, drawing methodological parallels to Sparse Sinkhorn Attention [2002.11296]. While S2T2 applies these mathematical operators for inter-vocabulary token translation, Sparse Sinkhorn Attention exploits them for sparse and efficient self-attention patterns via permutation of token blocks. Both fields share core theoretical machinery: Sinkhorn-Knopp normalization, doubly-stochastic constraints, and learned sparsity, but differ in application—token alignment versus attention routing. This broader context situates S2T2 as part of a trend leveraging optimal transport and differentiable combinatorial optimization techniques for scalable, sparse, and adaptable neural architectures [2002.11296, 2411.00593].

Source: https://www.emergentmind.com/topics/sparse-sinkhorn-token-translation-s2t2