Aleph Alpha GmbH

Engineering Blog 21. Apr. 2020

Schnelles Tokenizing natürlicher Sprache Tokenizing Natural Language, Quickly

Wir bei Aleph Alpha versuchen, den Stand der Technik des Maschinellen Lernens voranzutreiben. Wir setzen dazu auf die Programmiersprache Rust, daher war, als es um das Tokenizing ging, huggingface's Tokenizers crate die logische Wahl. Wir haben uns mit einem pull request bedankt, der die Tokenizer zwischen Threads verschiebbar macht.

Unser Vice President of Engineering, Andre 'llogiq' Bogus fand heraus, dass obwohl tokenizers dem neuesten Stand der Technik entspricht, die Zerlegung nicht so schnell wie technisch möglich läuft.

Also tat er sich mit Andrew Gallant (in Rust-Kreisen als 'BurntSushi' bekannt), um eine Zerlegungsmethode zu schaffen, die genau wie huggingface's wordpiece tokenizer funktioniert, aber schneller.

Die zugrunde liegende Erkenntnis ist, dass wordpiece vom Ende eines Wortes sucht, bis es ein Token gefunden hat. Dies macht es für Eingaben langsam, die lange Worte mit vielen Token enthalten. Auf Basis von BurntSushi's fst Bibliothek erstellten die beiden eine Methode, die ein Eingabewort linear durchsucht, und nur gelegentlich backtrackt, wenn es Token in der Eingabe gibt, die einen Präfix teilen.
At Aleph Alpha, we try to advance the state of the art of machine learning. We are also very bullish on the Rust programming language, so when it came to tokenizing, huggingface's tokenizers crate was the obvious solution for us. They even have a pull request from us to make their tokenizers `Send + Sync`.

However, while delving into the code, our VP of Engineering, Andre 'llogiq' Bogus found that despite being the state of the art, the tokenization wasn't as fast as it could be.

So he teamed up with Andrew Gallant (better known as 'BurntSushi') to create a tokenizing method that works exactly the same as huggingface's wordpiece tokenizer, but faster.

The insight here is that wordpiece scans from each word's end until it finds a token. This makes it slow if the input contains long words with many tokens. Using BurntSushi's fst crate, they created a method that does a linear scan, only occasionally resorting to backtracking when there are tokens in the input that share a prefix.

Andre Bogus


VP of Engineering

Durch die Verwendung je eines Finite State Transducers für Wortanfangs- und Folgetoken kann unser Tokenizer prefixes Byte-weise ermitteln, und bei jedem Byte entscheiden, ob dies ein komplettes Token ist, und ob das Vokabular noch längere Token mit demselben Präfix enthält. Das so minimierte Datenaufkommen ist ideal für moderne Prozessoren, bei denen jeder Speicherzugriff, der nicht in den Cache passt, zeitaufwändig ist.

In einem Benchmark zwischen wordpiece und aleph-alpha-tokenizer über einen kleinen Korpus deutscher Sätze fanden wir, dass der verbesserte Algorithmus Sätze mit langen Worten, die in viele Tokens zerlegt werden, am meisten beschleunigt (bis zu etwa einem Faktor Sechs). Sogar bei kurzen Sätzen mit ein-Token-Wörtern, der nachvollziehbar beste Fall für wordpiece, ist aleph-alpha-tokenizer um mindestens 20% schneller (wenn man ihn als Model benutzt, die selbe Schnittstelle, die auch wordpiece implementiert).

Abgesehen von der algorithmischen Verbesserung bemerkte llogiq, dass die internen Schnittstellen bei tokenizers für die Verwendung aus python oder javascript optimiert sind. Das ist sinnvoll, weil die Mehrzahl der Nutzer diese Sprachen verwenden. Allerdings bezahlt tokenizers diese Vereinfachung mit zusätzlichen Speicheranforderungen und Datenkopieen. Diese Kosten kann man vermeiden, wenn man direkt in Rust arbeitet.

Also baute er seinen Prototypen in eine komplette Bibliothek aus, die entweder alleine oder zusammen mit tokenizers verwendet werden kann. Die unabhängige Version ist weniger flexibel, was die Eingabetokens angeht, dafür einfacher zu nutzen. Dies wird sich allerdings in Zukunft ändern, wenn wir weitere Nutzungsfälle mit der Bibliothek abdecken. Außerdem erlaubt die Rust-Schnittstelle dem Nutzer, den Token ID Typ vorzugeben, was für uns die Umsetzung der IDs in andere Typen eliminiert, wenn die weitere Verarbeitung andere Datentypen als &[u32] verwendet.

Die Bibliothek hat noch keine Anbindung zur Verwendung aus anderen Sprachen. Wir planen, einen weiteren pull request an huggingface zu senden, um den aleph-alpha-tokenizer als optionales feature einzubinden.
By using one finite state transducer for word start and follower tokens, the tokenizer can look up prefixes byte-wise, at each byte deciding if this is a complete token and if there are longer tokens with the same prefix. The minimal data required for this task makes it very suitable for modern processors where memory access that does not fit in caches are at a premium.

Benchmarking both wordpiece and aleph-alpha-tokenizer on a small corpus of German sentences, we find that the algorithmic improvement speeds up sentences with long words that tokenize into many tokens the most (up to roughly a factor of six). Even on short sentences with one-token words, which is arguably the best case for wordpiece, aleph-alpha-tokenizer is at least 20% faster (when used as a `Model`, the same trait wordpiece implements).

Apart from the algorithmic improvements, llogiq noticed that the internal interfaces of tokenizers were optimized for use from dynamic languages to simplify writing the python and javascript bindings. This makes sense as the majority of users will access tokenizers from one the bindings. However, this trade-off involves some cost in memory allocation and copying that can be avoided when working directly in Rust.

Thus he fleshed out the initial proof of concept to a library crate that can be used either standalone or in conjunction with tokenizers. The standalone version is much less flexible regarding tokens, but easier to setup, however, this is likely to change in the future as we move more use cases to the new interface. Also the standalone version allows the user to specify the token ID type, which removes the need to transform the token IDs if the model does not take `u32`s.

The crate lacks bindings to other languages for now. We plan to send yet another PR to huggingface to include it in tokenizers as an optional feature.