Быстрая токенизация WordPiece: Линейный алгоритм для BERT и не только

Содержание

1. Введение

Токенизация — это фундаментальный этап предварительной обработки для почти всех задач обработки естественного языка (Natural Language Processing, NLP), от анализа тональности до машинного перевода. Современные модели, такие как BERT, GPT-3 и XLNet, полагаются на субсловную токенизацию, в частности на алгоритмы вроде WordPiece, чтобы разбивать текст на значимые единицы, одновременно обрабатывая слова вне словаря. В данной статье рассматривается критическое узкое место: эффективность самого алгоритма токенизации WordPiece.

Авторы отмечают, что стандартная жадная стратегия «сначала самое длинное совпадение» (MaxMatch), используемая в WordPiece, имеет неоптимальную временную сложность — $O(n^2)$ или $O(nk)$, где $n$ — длина входных данных, а $k$ — максимальная длина токена. Они предлагают LinMaxMatch — новый алгоритм, достигающий строгой линейной временной сложности $O(n)$ для токенизации отдельных слов, а также интегрированный алгоритм сквозной (End-to-End, E2E) токенизации WordPiece для произвольного текста.

2. Предпосылки и связанные работы

Токенизация WordPiece, используемая в BERT, включает два шага: (1) предварительную токенизацию текста на слова на основе пробелов и знаков препинания и (2) применение алгоритма MaxMatch к каждому слову с использованием фиксированного словаря. MaxMatch, также используемый в китайской сегментации слов, итеративно выбирает самый длинный префикс оставшейся строки, который соответствует токену из словаря.

Предыдущие эффективные алгоритмы были ограничены сложностью $O(n^2)$ или $O(nk)$. Фактор $k$ проблематичен, поскольку словари могут содержать длинные токены (например, «supercalifragilisticexpialidocious»). Вклад данной статьи заключается в полном устранении этого зависящего от словаря фактора.

3. Алгоритм LinMaxMatch

3.1 Основная идея и вдохновение

LinMaxMatch вдохновлён алгоритмом поиска строк Ахо-Корасик. Ключевая идея заключается в том, чтобы рассматривать словарь как набор шаблонов, а входное слово — как текст. Предварительно вычисляя вспомогательные структуры данных на префиксном дереве словаря, алгоритм может избежать возвратов (backtracking) во время сопоставления, которые являются источником квадратичной сложности в наивном MaxMatch.

3.2 Структуры данных: префиксное дерево, ссылки сбоя и всплывающие сбои

Алгоритм использует три основные структуры данных, построенные из словаря на этапе офлайн-предобработки:

• Префиксное дерево (Trie): Стандартное дерево префиксов, где каждый узел представляет символ, а пути от корня представляют токены словаря.
• Ссылка сбоя (аналогично Ахо-Корасик): Для узла, представляющего строку $s$, его ссылка сбоя указывает на узел, представляющий самый длинный собственный суффикс $s$, который также является префиксом некоторого токена словаря. Это позволяет алгоритму эффективно «переключаться при сбое».
• Всплывающий сбой (Failure Pop): Новое дополнение. При переходе от узла по его ссылке сбоя алгоритм может «извлечь» один или несколько токенов, которые были сопоставлены и заканчиваются в этом узле. Этот список предварительно вычисляется и сохраняется.

3.3 Пошаговое описание алгоритма

Для входного слова длины $n$ алгоритм обрабатывает символы слева направо за один проход. Он поддерживает текущий узел в префиксном дереве. Для каждого нового символа:

1. Попытаться перейти по соответствующему ребру из текущего узла.
2. Если успешно, перейти к дочернему узлу. Если этот узел является конечной точкой токена словаря, записать его как потенциальное совпадение.
3. Если соответствующего ребра не существует, перейти по ссылке сбоя текущего узла. Собрать все токены, хранящиеся в связанном списке «всплывающих сбоев». Затем повторить попытку сопоставления того же входного символа с нового узла.

Этот процесс гарантирует, что каждый символ обрабатывается ровно один раз, обеспечивая время работы $O(n)$.

4. Сквозная (End-to-End, E2E) токенизация WordPiece

Для токенизации полных предложений или документов авторы предлагают интегрировать этап предварительной токенизации (разделение по пробелам/знакам препинания) с LinMaxMatch в единый алгоритм. Этот алгоритм E2E WordPiece также работает за линейное время относительно общей длины текста, избегая накладных расходов отдельных этапов. Он умело чередует обнаружение границ с сопоставлением субслов в рамках одного обхода префиксного дерева.

5. Экспериментальные результаты и производительность

В статье представлены убедительные тесты, сравнивающие их реализацию с двумя широко используемыми библиотеками: HuggingFace Tokenizers и TensorFlow Text.

Результаты показывают, что теоретическое преимущество $O(n)$ преобразуется в значительный прирост производительности в реальных условиях, особенно для длинных текстов или в средах, чувствительных к задержкам, таких как мобильный вывод (inference).

6. Ключевые выводы и анализ

Основной вывод

Фунментальный прорыв статьи заключается не просто в более быстром токенизаторе; это формальное признание токенизации как задачи поиска строк, решаемой классической теорией автоматов. Переосмыслив MaxMatch WordPiece с наивного жадного поиска в вариант Ахо-Корасика, авторы преодолевают многолетний разрыв между теоретической информатикой и прикладной NLP-инженерией. Это напоминает то, как статья о CycleGAN переосмыслила перевод изображений как задачу циклически-согласованного состязательного обучения, создав новую парадигму.

Логическая последовательность

Аргументация логически безупречна: 1) Выявление узкого места (сложность $O(n^2)$/$O(nk)$). 2) Вдохновение проверенным оптимальным алгоритмом (Ахо-Корасик). 3) Адаптация его к специфическим ограничениям MaxMatch (необходимость выводить все сопоставленные токены, а не только факт существования). 4) Введение новых структур данных (Всплывающие сбои) для обработки этих ограничений. 5) Расширение решения на весь конвейер (E2E). Последовательность от проблемы к вдохновению, адаптации и обобщению является образцом прикладных исследований.

Сильные стороны и недостатки

Сильные стороны: Гарантия $O(n)$ теоретически элегантна и практически подтверждена. Стоимость предобработки — это единовременные накладные расходы. Влияние на мобильный/периферийный NLP может быть существенным, как отмечают авторы и подтверждают тенденции в исследованиях ИИ на устройствах от таких институтов, как Google AI.

Потенциальные недостатки/упущения: В статье вскользь упоминается, но не глубоко анализируется нагрузка на память для расширенного префиксного дерева (ссылки сбоя и всплывающие сбои). Для огромных словарей (например, многоязычных моделей) это может быть нетривиально. Кроме того, тесты, хотя и впечатляющие, проводятся против конкретных реализаций. Сравнение с другими оптимизированными вариантами MaxMatch (например, с использованием суффиксных массивов) укрепило бы утверждение.

Практические рекомендации

Для практиков NLP: Немедленно проанализируйте задержку вашей токенизации. Если она составляет более 1-2% от общего времени вывода, особенно для потоковых или мобильных приложений, внедрение алгоритма, подобного LinMaxMatch, является оптимизацией с высокой окупаемостью. Для исследователей: Эта работа открывает дверь. Могут ли аналогичные подходы на основе автоматов оптимизировать Byte-Pair Encoding (BPE) или SentencePiece? Основная идея — замена возвратов предварительно вычисленными переходами состояний — широко применима к другим жадным алгоритмам сегментации в NLP.

7. Технические детали и математическая формулировка

Основу LinMaxMatch можно формализовать. Пусть:

• $V$ = Словарь, множество строк.
• $T$ = Префиксное дерево, построенное из $V$.
• Для узла префиксного дерева $u$ пусть $L(u)$ — строка, образованная символами от корня до $u$.
• Функция сбоя $f(u)$: Узел $v$ такой, что $L(v)$ является самым длинным собственным суффиксом $L(u)$, где $v$ также является узлом в $T$.
• Всплывающий сбой $P(u)$: Множество токенов словаря $t \in V$, где $t$ является суффиксом $L(u)$ и существует узел $w$ на пути от корня до $f(u)$ (или самого $u$), где $L(w) = t$. Это предварительно вычисляется для всех узлов.

Токенизация входной строки $S[1..n]$ происходит путём поддержания текущего узла $u$ (изначально корень) и индекса $i$. Пока $i \leq n$:

1. Если существует дочерний узел $v$ для $u$ с меткой ребра $S[i]$, установить $u = v$, $i = i+1$. Если $u$ отмечает конец токена словаря, отметить это.
2. Иначе, вывести все токены из $P(u)$. Установить $u = f(u)$. Не увеличивать $i$. Если $u$ — корень и ни один дочерний узел не соответствует $S[i]$, тогда $i = i+1$ (обработка неизвестных символов).

Линейная сложность $O(n)$ этого цикла очевидна, так как $i$ либо увеличивается, либо $u$ следует по ссылкам сбоя, что можно амортизировать до константы на шаг.

8. Фреймворк анализа: Пример из практики

Сценарий: Оптимизация токенизации для голосового помощника реального времени на смартфоне. Текущий конвейер использует стандартный токенизатор WordPiece из популярной библиотеки.

Применение фреймворка:

1. Базовое измерение: Профилирование показывает, что токенизация занимает ~15 мс на запрос, что составляет 12% от общей задержки конвейера (125 мс). Цель — <100 мс.
2. Анализ узкого места: Запросы часто длинные (разговорные предложения). Фактор $O(nk)$ значителен, потому что словарь содержит технические составные слова.
3. Проектирование решения: Заменить токенизатор на LinMaxMatch. Единовременная стоимость построения расширенного префиксного дерева со ссылками/всплывающими сбоями приемлема во время развёртывания модели.
4. Ожидаемый результат: Исходя из заявленного в статье ускорения в 5-8 раз, время токенизации должно снизиться до ~2-3 мс, сократив общую задержку до ~112 мс, что является уверенным шагом к цели. Гарантия постоянного времени также улучшает предсказуемость задержки.

Этот пример из практики иллюстрирует, как теоретические свойства алгоритма напрямую решают практические ограничения в современных системах ИИ.

9. Перспективы применения и направления будущих исследований

Последствия эффективной линейной токенизации выходят за рамки BERT:

• ИИ на устройствах и на периферии: По мере того, как модели, подобные архитектурам GPT, дистиллируются для телефонов, каждая сэкономленная миллисекунда на предобработке суммируется. Эта работа является основополагающей для следующего поколения отзывчивых NLP-приложений с возможностью работы офлайн.
• Потоковая токенизация: Для перевода в реальном времени или субтитров в прямом эфире сложность $O(n)$ обеспечивает предсказуемую производительность. Будущие исследования могут изучить инкрементальную токенизацию для бесконечных потоков.
• Многоязычие и большие словари: Масштабирование до словарей с 1M+ токенов на 100+ языках. Необходимы исследования для оптимизации занимаемой памяти предвычисленных структур, возможно, с использованием сжатых префиксных деревьев или приближённых методов.
• Интеграция с архитектурой модели: Может ли токенизация быть частично изучена или объединена с первым слоем трансформера? «Динамическая токенизация» на основе контекста остаётся открытой проблемой, но эффективные базовые алгоритмы, подобные этому, являются необходимым условием для таких исследований.

10. Ссылки

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. (Приведено как пример смены парадигмы через переосмысление).