Sparse Sinkhorn Token Translation (S2T2)

• The paper introduces S2T2 that leverages entropic regularized optimal transport to learn a sparse mapping between target and source tokens, improving perplexity and compression.
• S2T2 formulates token translation via the Sinkhorn-Knopp algorithm with marginal constraints and sparsity penalties, enabling efficient domain adaptation.
• Empirical results reveal that S2T2 outperforms baselines in protein sequence modeling by achieving lower perplexity and enhanced data compression with transferable translation matrices.

Sparse Sinkhorn Token Translation (S2T2) is an adaptation technique for pre-trained LLMs 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 (Feng et al., 2024).

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 LLM $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| = v$0 (“soft” tokens), feed these through $|V_s| = v$1 to predict the next source token, and translate the predicted source-token distribution back to $|V_s| = v$2. This bidirectional translation is achieved with a learned matrix $|V_s| = v$3 interpreted as a transport plan between $|V_s| = v$4 and $|V_s| = v$5 (Feng et al., 2024).

2. Mathematical Foundations of Sparse Sinkhorn Token Translation

The translation matrix $|V_s| = v$6 is directly parameterized in $|V_s| = v$7 and constrained such that: $|V_s| = v$8 with frequency marginals $|V_s| = v$9 typically uniform ($L_s$0, $L_s$1).

The S2T2 objective jointly optimizes tokenizer parameters $L_s$2 and $L_s$3: $L_s$4 where $L_s$5 is the one-hot encoding of a target sequence, $L_s$6 applies a regularized Sinkhorn projection, $L_s$7 is the fixed LM probability, and $L_s$8 enforces sparsity (e.g., $L_s$9 or group-L1 penalty).

The entropic regularized optimal transport solution is: $V_s$0 with $V_s$1 as a cost/score matrix and $V_s$2 produced by Sinkhorn-Knopp iterations:

• Initialize $V_s$3 elementwise.
• Iterate: $V_s$4, $V_s$5.
• Form $V_s$6. After convergence, $V_s$7 satisfies prescribed marginals. Sparsity regularization is implemented by backpropagating through Sinkhorn steps into $V_s$8 and hence $V_s$9.

3. Optimization and Implementation Procedure

S2T2 parameters $T_t$0 are optimized end-to-end with $T_t$1 unrolled Sinkhorn iterations per forward pass. The workflow is:

1. Sample target-domain data $T_t$2, tokenize with $T_t$3 and encode as $T_t$4.
2. Compute Sinkhorn-projected transport $T_t$5.
3. Feed “soft” source tokens $T_t$6 through the fixed $T_t$7 to calculate next-token cross-entropy.
4. Add $T_t$8 to penalize translation density.
5. Backpropagate through Sinkhorn iterations to update $T_t$9 and tokenizer merges.
6. Hyperparameters: regularization weight $V_t$0 (0–1), Sinkhorn temperature $V_t$1–$V_t$2, $V_t$3 Sinkhorn iterations, batch size $V_t$4, sequence length $V_t$5, AdamW ($V_t$6, $V_t$7), cosine-annealed schedule with 20% warm-up (Feng et al., 2024).

4. Empirical Performance and Benchmarking

S2T2 is evaluated on adaptation from English to protein sequence modeling (UniRef50). A new BPE tokenizer $V_t$8 (vocab size $V_t$9) reduces token length by $|V_t| = u$01.82$|V_t| = u$1. Performance is measured via:

• Perplexity (“perp”): LM accuracy conditioned on token translations.
• Bits per byte (BpB): $|V_t| = u$2.

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 (Feng et al., 2024).

5. Transferability Across Model Scales

A notable property of S2T2 is the portability of the learned translation matrix $|V_t| = u$3 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 $|V_t| = u$4. Experiments show:

• Direct transfer and initialization with fixed $|V_t| = u$5 provides immediate perplexity and compression gains relative to both original and target-domain tokenizers.
• Learning $|V_t| = u$6 on a small model is approximately 7$|V_t| = u$7 cheaper than training the same on the larger model.
• Figure 1 in (Feng et al., 2024) shows that transferred S2T2 improves held-out loss from the outset, regardless of sparsity penalty $|V_t| = u$8.

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 (Feng et al., 2024).

Noted limitations and open questions include:

• Computational overhead of Sinkhorn iterations and need to carefully tune temperature $|V_t| = u$9 and number of iterations.
• The choice of the sparsity penalty $L_s$0 strongly affects the trade-off between translation matrix density and accuracy.
• Extension to multi-domain or multilingual adaptation would require block-structured, potentially hierarchical $L_s$1.
• 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 (Tay et al., 2020). 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 (Tay et al., 2020, Feng et al., 2024).