huggingface / huggingface/swift-transformers

Tokenizer.encode is quadratic in the number of added tokens (~29 ms per newline with Gemma-3)

Open Beginner friendly
#383 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
Swift
Stars
1.4k
Forks
209
Avg merge
2d 7h
Merged PRs (30d)
4

Description

# `Tokenizer.encode` is quadratic in the number of added tokens (~29 ms per newline with Gemma-3)

## Summary

`PreTrainedTokenizer` builds the added-token splitter as a single `NSRegularExpression` with **one capture group per added token** ([`Tokenizer.swift:517-523`](https://github.com/huggingface/swift-transformers/blob/main/Sources/Tokenizers/Tokenizer.swift#L517-L523)):

```swift
let addedTokensRegexString = unwrappedAddedTokens.map {
let token = NSRegularExpression.escapedPattern(for: $0.content)
let prefix = $0.prefix ? #"\s*"# : ""
let suffix = $0.suffix ? #"\s*"# : ""
return "\(prefix)(\(token))\(suffix)"
}.joined(separator: "|")
```

For a checkpoint with many added tokens this makes `encode` cost **~29 ms for every character that can begin an added token**. Gemma-3 has 6,415 added tokens, of which 6,322 start with `<` and 31 with `\n` — so any multi-line prompt pays ~29 ms per line.

Measured on `mlx-community/gemma-3-12b-it-4bit`, swift-transformers 1.3.3, release build, M-series Mac:

| input | `encode` |
|---|---|
| 44 chars, 1 line | 0.03 ms |
| same text ×10, newline-separated | 265 ms |
| same text ×40, newline-separated | **1148 ms** |
| single `"\n"` | 29.30 ms |
| single `"<"` | 28.88 ms |
| single `" "` | 0.004 ms |
| `"hello"` | 0.005 ms |

The same 4,175-character prompt costs **0.543 ms** via the Python `tokenizers` package on the same machine.

Cost tracks the *number of added tokens*, not the input, and is quadratic — truncating Gemma-3's `added_tokens` list and re-timing a single `"\n"`:

| added tokens | `encode("\n")` |
|---|---|
| 0 | 0.003 ms |
| 800 | 0.487 ms |
| 1,600 | 1.820 ms |
| 3,200 | 7.108 ms |
| 6,415 | 28.655 ms |

Doubling the added-token count quadruples the time.

## Cause

It is the capture groups, not the alternation. Reproduced with **no swift-transformers involved** — plain Foundation, 6,415 alternatives, matched against `"\n"`:

```swift
import Foundation

func toks(_ n: Int) -> [String] { (0.." } + ["\n"] }

func time(_ label: String, _ n: Int, capturing: Bool) {
let alts = toks(n).map { NSRegularExpression.escapedPattern(for: $0) }
.map { capturing ? "(\($0))" : $0 }
let re = try! NSRegularExpression(pattern: alts.joined(separator: "|"))
let text = "\n"
let r = NSRange(text.startIndex..`, ``, ``), unused tokens (``), newlines, tabs, the `▁` marker, bare `<`/`>` in prose, and plain text.

That is a targeted fix rather than a general one — a checkpoint that *does* set `lstrip`/`rstrip` on many added tokens would still hit the quadratic. A prefix automaton over the added tokens (as the Rust `tokenizers` crate uses) would remove it in all cases, if you'd prefer that direction.

## Environment

- swift-transformers 1.3.3 (`2fa33e1`), latest release at time of writing
- macOS 26 (Darwin 27.0.0), Apple silicon, release build (`-c release`)
- Reproduces via `AutoTokenizer.from(modelFolder:)` and `AutoTokenizer.from(pretrained:)` alike
- Not reproduced on `openai-community/gpt2`, `bert-base-uncased`, or `mlx-community/Qwen2.5-0.5B-Instruct-4bit` — all have few added tokens

I'm happy to open a PR with the patch above plus a regression test if that's useful.

Contributor guide

No contributing guide indexed for this repository

Research direction

Start in Sources/Tokenizers/Tokenizer.swift around lines 517-523, where the added-token regular expression is assembled, and read String+PreTokenization.swift around split(by:). Reproduce the slowdown with a large added-token list and newline input, then verify that the targeted regex change preserves token IDs while removing the capture-group overhead; add the mentioned regression test.

Written by the indexing model from the issue text.

Assessment

Tech stack
swift
Domain
performance
Issue type
Bug
Difficulty
2/5
Estimated time
1-3 hours
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
82/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.