Tokenisation Fast WordPiece : Un algorithme à temps linéaire pour BERT et au-delà

Table des matières

1. Introduction

La tokenisation est l'étape de prétraitement fondamentale pour presque toutes les tâches de Traitement Automatique des Langues (TAL), de l'analyse de sentiments à la traduction automatique. Les modèles modernes comme BERT, GPT-3 et XLNet reposent sur la tokenisation par sous-mots, en particulier des algorithmes comme WordPiece, pour décomposer le texte en unités significatives tout en gérant les mots hors vocabulaire. Cet article aborde un goulot d'étranglement critique : l'efficacité de l'algorithme de tokenisation WordPiece lui-même.

Les auteurs identifient que la stratégie gloutonne standard de correspondance la plus longue en premier (MaxMatch) utilisée dans WordPiece a une complexité temporelle sous-optimale — $O(n^2)$ ou $O(nk)$, où $n$ est la longueur de l'entrée et $k$ la longueur maximale d'un token. Ils proposent LinMaxMatch, un nouvel algorithme atteignant une complexité temporelle linéaire stricte de $O(n)$ pour la tokenisation d'un mot unique, et un algorithme intégré de Tokenisation WordPiece de bout en bout (E2E) pour le texte général.

2. Contexte et travaux connexes

La tokenisation WordPiece, telle qu'utilisée dans BERT, implique deux étapes : (1) la prétokenisation du texte en mots basée sur les espaces et la ponctuation, et (2) l'application de l'algorithme MaxMatch à chaque mot par rapport à un vocabulaire fixe. MaxMatch, également utilisé en segmentation de mots chinois, sélectionne itérativement le plus long préfixe de la chaîne restante qui correspond à un token du vocabulaire.

Les algorithmes efficaces antérieurs étaient limités à $O(n^2)$ ou $O(nk)$. Le facteur $k$ est problématique car les vocabulaires peuvent contenir des tokens longs (par exemple, "supercalifragilisticexpialidocious"). La contribution de cet article réside dans l'élimination complète de ce facteur dépendant du vocabulaire.

3. L'algorithme LinMaxMatch

3.1 Idée centrale et inspiration

LinMaxMatch s'inspire de l'algorithme de correspondance de chaînes Aho-Corasick. L'idée clé est de traiter le vocabulaire comme un ensemble de motifs et le mot d'entrée comme le texte. En précalculant des structures de données auxiliaires sur le trie du vocabulaire, l'algorithme peut éviter le retour en arrière pendant la correspondance, qui est la source de la complexité quadratique dans l'approche naïve de MaxMatch.

3.2 Structures de données : Trie, liens d'échec et "Failure Pops"

L'algorithme utilise trois structures de données centrales construites à partir du vocabulaire lors d'une phase de prétraitement hors ligne :

• Trie : Un arbre de préfixes standard où chaque nœud représente un caractère et les chemins depuis la racine représentent les tokens du vocabulaire.
• Lien d'échec (similaire à Aho-Corasick) : Pour un nœud représentant la chaîne $s$, son lien d'échec pointe vers le nœud représentant le plus long suffixe propre de $s$ qui est également un préfixe d'un token du vocabulaire. Cela permet à l'algorithme de "basculer" efficacement.
• "Failure Pop" : Une nouveauté. Lors du parcours depuis un nœud via son lien d'échec, l'algorithme peut "faire apparaître" un ou plusieurs tokens qui ont été appariés et se terminent à ce nœud. Cette liste est précalculée et stockée.

3.3 Explication détaillée de l'algorithme

Pour un mot d'entrée de longueur $n$, l'algorithme traite les caractères de gauche à droite en une seule passe. Il maintient un nœud courant dans le trie. Pour chaque nouveau caractère :

1. Tenter de suivre l'arête correspondante depuis le nœud courant.
2. En cas de succès, se déplacer vers le nœud enfant. Si ce nœud est un point de terminaison de token du vocabulaire, l'enregistrer comme une correspondance potentielle.
3. Si aucune arête correspondante n'existe, suivre le lien d'échec du nœud courant. Collecter tous les tokens stockés dans la liste "failure pop" associée. Puis, réessayer de faire correspondre le même caractère d'entrée depuis le nouveau nœud.

Ce processus garantit que chaque caractère est traité exactement une fois, assurant un temps $O(n)$.

4. Tokenisation WordPiece de bout en bout (E2E)

Pour tokeniser des phrases ou des documents complets, les auteurs proposent d'intégrer l'étape de prétokenisation (division sur les espaces/ponctuation) avec LinMaxMatch en un seul algorithme unifié. Cet algorithme E2E WordPiece fonctionne également en temps linéaire par rapport à la longueur totale du texte, évitant la surcharge des étapes séparées. Il entrelace astucieusement la détection des limites avec la correspondance des sous-mots dans le même cadre de parcours de trie.

5. Résultats expérimentaux et performances

L'article présente des benchmarks convaincants comparant leur implémentation à deux bibliothèques largement utilisées : HuggingFace Tokenizers et TensorFlow Text.

Les résultats démontrent que l'avantage théorique $O(n)$ se traduit par des gains de performances significatifs dans le monde réel, en particulier pour les textes longs ou dans des environnements sensibles à la latence comme l'inférence mobile.

6. Principales observations et analyse

Observation fondamentale

La percée fondamentale de l'article n'est pas seulement un tokeniseur plus rapide ; c'est la reconnaissance formelle de la tokenisation comme un problème de correspondance de chaînes résoluble par la théorie classique des automates. En reformulant le MaxMatch de WordPiece d'une recherche gloutonne naïve en une variante d'Aho-Corasick, les auteurs comblent un fossé vieux de plusieurs décennies entre l'informatique théorique et l'ingénierie appliquée du TAL. Cela rappelle la manière dont l'article CycleGAN a reformulé la traduction d'images comme un problème d'apprentissage antagoniste à cohérence cyclique, créant un nouveau paradigme.

Flux logique

L'argumentation est logiquement irréprochable : 1) Identifier le goulot d'étranglement (complexité $O(n^2)$/$O(nk)$). 2) S'inspirer d'un algorithme optimal et éprouvé (Aho-Corasick). 3) L'adapter aux contraintes spécifiques de MaxMatch (besoin de sortir tous les tokens appariés, pas seulement l'existence). 4) Introduire de nouvelles structures de données (Failure Pops) pour gérer ces contraintes. 5) Étendre la solution au pipeline complet (E2E). Le flux allant du problème à l'inspiration, puis à l'adaptation et à la généralisation, est un exemple de recherche appliquée exemplaire.

Points forts et limites

Points forts : La garantie $O(n)$ est élégante théoriquement et vérifiée pratiquement. Le coût de prétraitement est une surcharge unique. L'impact sur le TAL mobile/edge pourrait être substantiel, comme le notent les auteurs et comme le soutiennent les tendances de la recherche en IA sur appareil d'institutions comme Google AI.

Limites/omissions potentielles : L'article aborde légèrement mais n'analyse pas en profondeur la surcharge mémoire pour le trie amélioré (liens d'échec et pops). Pour les vocabulaires massifs (par exemple, les modèles multilingues), cela pourrait être non négligeable. De plus, le benchmark, bien qu'impressionnant, est réalisé contre des implémentations spécifiques. Une comparaison avec d'autres variantes optimisées de MaxMatch (par exemple, utilisant des tableaux de suffixes) renforcerait l'affirmation.

Perspectives exploitables

Pour les praticiens du TAL : Profilez immédiatement la latence de votre tokenisation. Si elle représente plus de 1 à 2 % de votre temps d'inférence total, en particulier pour les applications de streaming ou mobiles, l'adoption d'un algorithme comme LinMaxMatch est une optimisation à fort retour sur investissement. Pour les chercheurs : ce travail ouvre la porte. Des approches similaires basées sur les automates peuvent-elles optimiser l'encodage par paires d'octets (BPE) ou SentencePiece ? L'idée centrale — remplacer le retour en arrière par des transitions d'état précalculées — est largement applicable à d'autres algorithmes de segmentation gloutons en TAL.

7. Détails techniques et formulation mathématique

Le cœur de LinMaxMatch peut être formalisé. Soit :

• $V$ = Vocabulaire, un ensemble de chaînes.
• $T$ = Trie construit à partir de $V$.
• Pour un nœud de trie $u$, soit $L(u)$ la chaîne formée par les caractères de la racine à $u$.
• Fonction d'échec $f(u)$ : Le nœud $v$ tel que $L(v)$ est le plus long suffixe propre de $L(u)$ où $v$ est également un nœud dans $T$.
• "Failure Pop" $P(u)$ : L'ensemble des tokens du vocabulaire $t \in V$ où $t$ est un suffixe de $L(u)$ et il existe un nœud $w$ sur le chemin de la racine à $f(u)$ (ou $u$ lui-même) où $L(w) = t$. Ceci est précalculé pour tous les nœuds.

La tokenisation de la chaîne d'entrée $S[1..n]$ procède en maintenant un nœud courant $u$ (initialement la racine) et un index $i$. Tant que $i \leq n$ :

1. S'il existe un enfant $v$ de $u$ avec l'étiquette d'arête $S[i]$, définir $u = v$, $i = i+1$. Si $u$ marque la fin d'un token du vocabulaire, le noter.
2. Sinon, sortir tous les tokens dans $P(u)$. Définir $u = f(u)$. Ne pas incrémenter $i$. Si $u$ est la racine et qu'aucun enfant ne correspond à $S[i]$, alors $i = i+1$ (gérer les caractères inconnus).

La complexité $O(n)$ de cette boucle est claire, car $i$ augmente soit, soit $u$ suit des liens d'échec, ce qui peut être amorti à une constante par étape.

8. Cadre d'analyse : une étude de cas

Scénario : Optimisation de la tokenisation pour un assistant vocal en temps réel sur un smartphone. Le pipeline actuel utilise un tokeniseur WordPiece standard d'une bibliothèque populaire.

Application du cadre :

1. Mesure de référence : Le profilage montre que la tokenisation consomme ~15 ms par requête, soit 12 % de la latence totale du pipeline (125 ms). L'objectif est <100 ms.
2. Analyse du goulot d'étranglement : Les requêtes sont souvent longues (phrases parlées). Le facteur $O(nk)$ est significatif car le vocabulaire contient des mots composés techniques.
3. Conception de la solution : Remplacer le tokeniseur par LinMaxMatch. Le coût unique de construction du trie amélioré avec les liens d'échec/pops est acceptable lors du déploiement du modèle.
4. Résultat attendu : Sur la base de l'accélération de 5 à 8x de l'article, le temps de tokenisation devrait chuter à ~2-3 ms, réduisant la latence totale à ~112 ms, un pas solide vers l'objectif. La garantie de temps constant améliore également la prévisibilité de la latence.

Cette étude de cas illustre comment les propriétés théoriques de l'algorithme répondent directement aux contraintes pratiques des systèmes d'IA modernes.

9. Perspectives d'application et orientations futures

Les implications d'une tokenisation efficace à temps linéaire s'étendent au-delà de BERT :

• IA sur appareil et en périphérie : Alors que des modèles comme les architectures de type GPT sont distillés pour les téléphones, chaque milliseconde économisée dans le prétraitement se cumule. Ce travail est fondamental pour la prochaine génération d'applications TAL réactives et fonctionnant hors ligne.
• Tokenisation en flux continu : Pour la traduction en temps réel ou le sous-titrage en direct, la complexité $O(n)$ garantit des performances prévisibles. Les travaux futurs pourraient explorer la tokenisation incrémentielle pour des flux infinis.
• Vocabulaires multilingues et volumineux : Passage à l'échelle vers des vocabulaires avec 1M+ de tokens dans 100+ langues. Des recherches sont nécessaires pour optimiser l'empreinte mémoire des structures précalculées, peut-être en utilisant des tries compressés ou des méthodes approximatives.
• Intégration avec l'architecture du modèle : La tokenisation pourrait-elle être partiellement apprise ou fusionnée avec la première couche d'un transformateur ? La "Tokenisation Dynamique" basée sur le contexte reste un défi ouvert, mais des algorithmes de base efficaces comme celui-ci sont un prérequis pour une telle exploration.

10. Références

1. Devlin, J., Chang, M. W., Lee, K., & Toutanova, K. (2019). BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. NAACL-HLT.
2. Aho, A. V., & Corasick, M. J. (1975). Efficient string matching: an aid to bibliographic search. Communications of the ACM.
3. Kudo, T. (2018). Subword Regularization: Improving Neural Network Translation Models with Multiple Subword Candidates. ACL.
4. Wolf, T., et al. (2020). HuggingFace's Transformers: State-of-the-art Natural Language Processing. arXiv preprint arXiv:1910.03771.
5. Google. (2020). TensorFlow Text. https://www.tensorflow.org/text
6. Brown, T. B., et al. (2020). Language Models are Few-Shot Learners. NeurIPS.
7. Zhu, J.-Y., Park, T., Isola, P., & Efros, A. A. (2017). Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks. ICCV. (Cité comme exemple de reformulation changeant de paradigme).