Детерминированное построение таблицы слияний BPE

Один токен для повторяющихся байтов

В главе 1 мы различали единицы текста и ID, которыми они обозначаются. В таблице целых слов встретившееся слово может получить один ID, но конечному списку написаний нужен запасной вариант для незнакомого слова. Если начать со всех возможных байтов, это ограничение исчезает: любой документ UTF-8 уже состоит из этих байтов. Однако повторяющаяся последовательность байтов по-прежнему занимает несколько позиций в последовательности токенов. Кодирование пар байтов (BPE) учится по таким повторам строить более крупные единицы для повторного использования.

Здесь ID токена обозначает запись словаря, а последовательность байтов токена — точные байты, которые эта запись представляет, с сохранением их порядка. Изначально ID от 0 до 255 представляют по одному соответствующему байту. ID служит меткой; это не кодовая точка Unicode и не численная оценка смысла. Добавление новых записей по данным — обучение токенизатора. Нейронные веса в этой главе не оптимизируются.

В главе 2 роли документов были зафиксированы до обучения. Алгоритм читает только обучающие документы и сохраняет для каждого отдельную последовательность токенов. Полностью разберём обучение на двух коротких документах:

train-aaa = aaa = [97, 97, 97]
train-aba = aba = [97, 98, 97]

В ASCII байт a имеет десятичное значение 97, или 61 в шестнадцатеричной записи; байт b — 98, или 62. На первом раунде считаются все пары соседних позиций токенов. В aaa пара (97,97) находится в позициях 0–1 и 1–2. Средний токен входит в оба окна, поэтому учитываются два вхождения кандидата. Отдельный документ aba добавляет одно (97,98) и одно (98,97):

Наибольшее число вхождений у (97,97). Это первое полученное правило имеет ранг слияния ноль — позицию ноль в упорядоченной таблице правил. Ему назначается новый ID 256 с байтами 61 61.

Затем замена один раз проходит по каждому документу слева направо. В aaa совпадение в позициях 0–1 использует оба входных токена, и вместо них в результат записывается 256. Просмотр продолжается с позиции 2: оставшийся 97 копируется. Средний входной токен уже использован и не может участвовать в следующей замене. В aba совпадения с этим правилом нет, поэтому документ не меняется:

train-aaa = [256, 97]
train-aba = [97, 98, 97]

Число вхождений пары-кандидата равно двум, а число замен — одному. Первое число нужно для выбора правила по всем соседним окнам; второе считает совпадения, фактически использованные при замене. При перекрывающихся вхождениях одного числа недостаточно, чтобы определить другое.

Второй раунд заново считает пары в текущих последовательностях. Теперь (97,98), (98,97) и (256,97) встречаются по одному разу. Прежнее число вхождений (97,97) больше не используется. При равных частотах в этом курсе сначала сравниваются левые числовые ID, затем правые. Выигрывает (97,98): его левый ID 97 меньше 98 и 256.

Правило ранга один получает новый ID 257 с байтами 61 62. Его единственное совпадение — первые два токена aba. Получаются последовательности:

train-aaa = [256, 97]
train-aba = [257, 97]

Правило добавляется и при одном вхождении; алгоритм не требует, чтобы пара встретилась хотя бы дважды. Общее число позиций токенов в двух документах уменьшается с шести до пяти, затем до четырёх. Если заменить ID в [256,97] их байтами, получится 61 61 61; для [257,97] получится 61 62 61. Исходные байты каждого документа сохранены. Словарь при этом вырос с 256 до 258 записей. Число записей словаря, число позиций токенов и длина в байтах — разные величины.

Граница документа завершает последовательность кандидатов. В двух однобайтовых документах a | a черта изображает границу и не входит во входные данные. Ни в одном документе нет двух позиций, поэтому кандидата (97,97) нет. Внутри документа пробел — обычный байт; отдельных правил для границ слов или символов у этого алгоритма нет.

Общее правило для каждого раунда

В примере повторяется один порядок действий: посчитать текущие соседние окна внутри каждого документа, выбрать пару, назначить объединённым байтам новый ID и заменить эту пару за один проход по каждому документу. Повторный подсчёт начинается только после завершения этих проходов. Выбор пары и новую запись обозначает формула:

Оценка состоит из трёх компонентов и сравнивается лексикографически: результат определяет первый различающийся компонент. Сначала максимизируется число вхождений $C(a,b)$. При равных частотах максимум $-a$ соответствует меньшему левому числовому ID. Если равны и левые ID, максимум $-b$ выбирает меньший правый ID. Поэтому оценки второго раунда $(1,-97,-98)$, $(1,-98,-97)$ и $(1,-256,-97)$ выбирают пару $(97,98)$. Это правило для равных частот явно принято в курсе.

Звёздочки отмечают выбранную пару. Во втором выражении объединение $a^{*}$ и $b^{*}$ означает приписывание байтов правого токена к байтам левого, чтобы задать новую запись $m^{*}$. Например, токен 256 представляет 61 61, а токен 97 — 61; последующее правило для (256,97) сохранило бы 61 61 61. Числовые ID при этом не складываются и не склеиваются как десятичные строки "256" и "97".

Код на Rust задаёт те же приоритеты без смены знака у беззнаковых ID. windows(2) даёт перекрывающиеся окна-кандидаты, числовое упорядочение пар ставит меньшую пару первой, а выбранный кандидат меняется только при строго большем числе вхождений. Отдельный индекс замены использует два входных токена при совпадении и один в остальных случаях. Новые токены записываются в результат текущего раунда и в том же проходе заново не просматриваются.

Запрошенное число слияний — верхняя граница. Если кандидатов нет, обучение заканчивается раньше; в остальных случаях даже одно вхождение победителя позволяет добавить правило. При запросе нуля раундов исходные последовательности и словарь из 256 однобайтовых записей остаются без изменений.

Обозначения в формуле

Начальным байтовым записям соответствуют ID 0..=255. Если $r$ — ранг правила при нумерации с нуля, новый ID равен $256+r$. После $k$ успешных раундов словарь содержит $256+k$ записей, включая все 256 исходных байтов. Эти выражения описывают число и нумерацию записей словаря, а не число позиций токенов в документе.

От словаря целых слов к повторному слиянию пар

Таблица целых слов представляет известное слово одной единицей поиска. Короткий пример на Rust строит словарь по словам из low lower и new newest, разделённым пробелами. Четыре написания сортируются; lower получает ID 2. Для lowering и lowest записей нет, поэтому оба возвращают один ID неизвестного токена 0. При таком поиске различие между этими написаниями теряется.

У примера с целыми словами своя нумерация ID. Его неизвестный ID 0 не связан с ID 0 байтового алгоритма обучения: последний представляет один байт 00.

pub fn fit_whole_word_vocabulary(documents: &[&str]) -> BTreeMap<String, u32> {
    let words = documents
        .iter()
        .flat_map(|document| document.split_whitespace())
        .collect::<BTreeSet<_>>();
    words
        .into_iter()
        .enumerate()
        .map(|(index, word)| (word.to_owned(), index as u32 + 1))
        .collect()
}

/// Returns the fitted ID or one shared unknown-word bucket.
pub fn whole_word_id(vocabulary: &BTreeMap<String, u32>, word: &str) -> u32 {
    vocabulary.get(word).copied().unwrap_or(UNKNOWN_WORD_ID)
}

В статье Филипа Гейджа 1994 года «A New Algorithm for Data Compression» частые соседние пары байтов многократно заменялись неиспользуемыми байтами, а таблица замен сохранялась для восстановления исходной последовательности. С нашей задачей это связывает идея одного символа для повторяющейся пары. Набор свободных байтов у Гейджа был конечным; его коды и процедура сжатия отличаются от новых ID типа u32 и полного повторного подсчёта в каждом следующем раунде нашего алгоритма.

Рико Сеннрих, Барри Хэддоу и Александра Бёрч адаптировали повторное слияние пар для редких слов в нейронном машинном переводе. В разделе 3.2 их статьи 2016 года о подсловных единицах исходными единицами служат символы; слова представлены отдельно с маркером конца слова, а последовательности символов объединяются внутри слова. Благодаря этому разные слова могут использовать общие полученные единицы, вместо того чтобы каждому требовалась отдельная запись целого слова.

Вариант курса начинается со всех 256 байтов и ограничивает пары границами документов. Полученный токен может включать пробел или заканчиваться внутри символа UTF-8. Ни один из исторических источников не задаёт наше правило числового выбора при равной частоте; слияние пар также не гарантирует, что токен окажется языковой морфемой. Пример с целыми словами показывает ограничение поиска в словаре. Он не сравнивает качество нейронных моделей и пока не кодирует незнакомый текст с помощью BPE.

Те же операции в существующем коде на Rust

Существующий счётчик пар обходит вектор токенов каждого документа. Вызов windows(2) включает оба окна из aaa; новый обход для каждого документа исключает пары через границу документов. TokenPair упорядочивает кандидатов сначала по левому числовому ID, затем по правому. BTreeMap обходит их в этом порядке, а условие count > best_count меняет выбранную пару только при большем числе вхождений. При равенстве остаётся наименьшая пара. Так реализуются три приоритета формулы выбора:

pub fn count_adjacent_pairs(sequences: &[Vec<u32>]) -> BTreeMap<TokenPair, usize> {
    let mut counts = BTreeMap::new();
    for sequence in sequences {
        for window in sequence.windows(2) {
            let pair = TokenPair::new(window[0], window[1]);
            *counts.entry(pair).or_insert(0) += 1;
        }
    }
    counts
}

/// Selects the greatest count, breaking ties by the smallest `(left, right)` IDs.
pub fn choose_most_frequent_pair(
    counts: &BTreeMap<TokenPair, usize>,
) -> Option<(TokenPair, usize)> {
    let mut winner = None;
    for (&pair, &count) in counts {
        match winner {
            None => winner = Some((pair, count)),
            Some((_, best_count)) if count > best_count => winner = Some((pair, count)),
            Some(_) => {}
        }
    }
    winner
}

Функция замены читает вход по индексу и отдельно записывает результат. При совпадении она записывает новый ID и сдвигает индекс на два; иначе копирует один ID и сдвигает индекс на один. Индекс не возвращается к использованным токенам, а результат не читается как новый вход. Поэтому замены в пределах одного прохода не перекрываются:

pub fn replace_pair_left_to_right(
    sequence: &[u32],
    pair: TokenPair,
    replacement: u32,
) -> (Vec<u32>, usize) {
    let mut output = Vec::with_capacity(sequence.len());
    let mut replacements = 0;
    let mut index = 0;

    while index < sequence.len() {
        if index + 1 < sequence.len()
            && sequence[index] == pair.left
            && sequence[index + 1] == pair.right
        {
            output.push(replacement);
            replacements += 1;
            index += 2;
        } else {
            output.push(sequence[index]);
            index += 1;
        }
    }

    (output, replacements)
}

BpeTrainer::train принимает проверенное разбиение CorpusPartitions из главы 2. Он получает training_documents(), сохраняет стабильные ID этих документов и преобразует байты их содержимого в отдельные начальные векторы токенов. Валидационные и тестовые документы в этом процессе не дают частот пар. Обучающая выборка задаёт зависимость алгоритма от данных; она не запрещает другому коду читать отложенные документы:

    pub fn train(self, partitions: &CorpusPartitions<'_>) -> Result<BpeTraining, BpeTrainingError> {
        let available_merge_ids = u128::from(u32::MAX) - u128::from(u8::MAX);
        if self.max_merges as u128 > available_merge_ids {
            return Err(BpeTrainingError::new(
                "requested merge count exceeds the u32 token-ID space",
            ));
        }
        let training_documents = partitions.training_documents();
        let document_ids = training_documents
            .iter()
            .map(|document| document.id().to_owned())
            .collect::<Vec<_>>();
        let sequences = training_documents
            .iter()
            .map(|document| bytes_to_tokens(document.text().as_bytes()))
            .collect::<Vec<_>>();

        learn_from_token_sequences(self.max_merges, document_ids, sequences)
    }

Внутренний цикл обучения начинает с 256 однобайтовых последовательностей. В каждом успешном раунде текущий размер словаря становится новым ID. Алгоритм объединяет байты выбранных записей, заменяет пару отдельно в каждом документе и сохраняет правило с последовательностью байтов. На следующей итерации пары считаются заново по полученным последовательностям. В каждом правиле записываются и число вхождений кандидата, и число замен, поэтому различие этих величин можно проверить.

Демонстрационная программа выполняет два отдельных примера. Сначала она запрашивает восемь слияний по восьми обучающим документам главы 2; два валидационных и два тестовых документа остаются отложенными. Затем программа выполняет два раунда для aaa/aba. Эта короткая трассировка не является первыми двумя рангами обучения на корпусе:

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let corpus = Corpus::from_json(CORPUS_JSON)?;
    let manifest = SplitManifest::from_json(SPLIT_MANIFEST)?;
    let partitions = manifest.partition(&corpus)?;
    let training = BpeTrainer::new(8).train(&partitions)?;

    println!("corpus checksum: {}", corpus.checksum());
    println!("statistics source: train only");
    println!("training documents: {:?}", training.training_document_ids());
    println!(
        "held out from trainer: validation={} test={}",
        partitions.documents(Partition::Validation).len(),
        partitions.documents(Partition::Test).len()
    );
    println!(
        "merge rounds: requested={} learned={}",
        training.requested_merges(),
        training.rules().len()
    );
    for rule in training.rules() {
        println!(
            "corpus rank {}: pair={} count={} replacements={} token={} bytes={:02x?}",
            rule.rank(),
            format_pair(rule.pair()),
            rule.candidate_count(),
            rule.replacement_count(),
            rule.token_id(),
            training
                .token_bytes(rule.token_id())
                .expect("learned token has bytes")
        );
    }

    let mut fixture = vec![bytes_to_tokens(b"aaa"), bytes_to_tokens(b"aba")];
    let mut fixture_vocabulary = (u8::MIN..=u8::MAX)
        .map(|byte| vec![byte])
        .collect::<Vec<_>>();
    println!("TRACE bpe-merges-v1 BEGIN");
    print_trace_stage(0, &fixture);
    for rank in 0..2 {
        let counts = count_adjacent_pairs(&fixture);
        let (winner, count) = choose_most_frequent_pair(&counts).expect("fixture has a pair");
        let token_id = BYTE_TOKEN_COUNT + rank;
        for (pair, candidate_count) in &counts {
            let selected = if *pair == winner { "yes" } else { "no" };
            println!(
                "CANDIDATE rank={rank} left={} right={} count={candidate_count} winner={selected}",
                pair.left(),
                pair.right()
            );
        }
        let mut merged_bytes = fixture_vocabulary[winner.left() as usize].clone();
        merged_bytes.extend_from_slice(&fixture_vocabulary[winner.right() as usize]);
        let mut replacements = 0;
        for sequence in &mut fixture {
            let (next, replaced) = replace_pair_left_to_right(sequence, winner, token_id);
            *sequence = next;
            replacements += replaced;
        }
        println!(
            "MERGE rank={rank} left={} right={} count={count} replacements={replacements} token={token_id} bytes_hex={}",
            winner.left(),
            winner.right(),
            format_hex(&merged_bytes)
        );
        fixture_vocabulary.push(merged_bytes);
        print_trace_stage(rank as usize + 1, &fixture);
    }
    println!("TRACE bpe-merges-v1 END");
    println!(
        "document barrier candidates for A=\"a\" B=\"a\": {}",
        count_adjacent_pairs(&[bytes_to_tokens(b"a"), bytes_to_tokens(b"a")]).len()
    );

    let words = fit_whole_word_vocabulary(&["low lower", "new newest"]);
    println!("historical whole-word types: {}", words.len());
    println!(
        "historical lookup lower: {}",
        whole_word_id(&words, "lower")
    );
    println!(
        "historical lookup lowering: {} (unknown)",
        whole_word_id(&words, "lowering")
    );
    println!("chapter 4 handoff: freeze ranks and encode arbitrary bytes");

    Ok(())
}

На корпусе правило ранга ноль выбирает (32,208): число вхождений кандидата равно 81, число замен — 81, новый ID при обучении — 256, а байты в шестнадцатеричной записи — 20 d0. Здесь 20 — пробел, а d0 начинает двухбайтовую последовательность UTF-8, продолжение которой в этот токен не вошло. Поэтому эта последовательность сама по себе не является корректной строкой UTF-8. Это следует из правил байтовых последовательностей в разделах 3–4 RFC 3629.

Один токен не обязан быть одним символом, одним словом или даже законченной строкой UTF-8. Алгоритм сохраняет байты и не подставляет выдуманный символ для такого фрагмента. Если объединить последовательности байтов всех токенов полного документа по порядку, исходные байты документа сохранятся.

Два раунда на одной схеме

На каждом этапе есть строка для train-aaa и строка для train-aba. Граница документов не позволяет концу одной строки образовать пару с началом другой. Число вхождений 2 в первой таблице относится к перекрывающимся окнам в aaa, а число замен 1 — к одному использованному совпадению. На следующем этапе показана изменённая последовательность, по которой вторая таблица заново считает пары.

Во второй таблице три кандидата встречаются по одному разу. Выбрана строка (97,98): равенство частот разрешается числовым порядком пар. Последний этап показывает [256,97] и [257,97], связывая две выполненные замены с сохранившимися отдельными документами. Выбор отмечен словами, звёздочкой и рамкой.

Необязательное воспроизведение и практика

Чтобы воспроизвести приведённые результаты, понадобится локальная копия репозитория, зафиксированная в нём версия инструментов Rust и зависимости из файла блокировки. Выполните команды из корня репозитория:

cargo test --locked -p llm-from-scratch --lib tokenizer::bpe_trainer::tests::
cargo test --locked -p ch03-learn-bpe-merges
cargo run --quiet --locked -p ch03-learn-bpe-merges | diff -u rust/demos/ch03-learn-bpe-merges/expected.txt -

Первая команда запускает шесть тестов алгоритма обучения, вторая — один тест библиотеки демонстрационного пакета. При успехе тесты сообщают об отсутствии ошибок. Последняя команда сравнивает стандартный вывод программы с ожидаемым выводом из репозитория. Если они совпадают, сравнение ничего не печатает и завершается с кодом ноль. В выводе есть и восемь правил по корпусу, и отдельная короткая трассировка, разобранная выше.

Разберите следствия правила, сверяясь с готовыми ответами:

1. Для одного документа [97,97,97,97] сопоставьте число соседних окон с числом совпадений, использованных за один проход замены слева направо.
2. Рассмотрите три тройки оценки после ранга ноль в коротком примере и объясните выбор (97,98).
3. Объясните отсутствие кандидата (97,98) у отдельных документов [97] | [98]. Здесь | изображает границу документов и не является входными данными.
4. Свяжите размер словаря с девятью успешными раундами из двенадцати запрошенных. Также укажите, что сохраняется при запросе нуля раундов.
5. Объясните, почему изменение числа повторов в валидационном документе не даёт этому алгоритму новой статистики пар. Достаточно рассуждения; редактировать корпус не нужно.
6. Объясните, какие байты получит гипотетическое правило (256,97), если токен 256 хранит 61 61, а токен 97 — 61.
7. Объясните, почему правило по корпусу с байтами 20 d0 нельзя напечатать как отдельный символ и почему это не делает некорректным весь документ.

От полученного словаря к модели

Результат обучения — упорядоченная таблица правил, последовательность байтов для каждой записи словаря и преобразованные обучающие последовательности. В коротком примере позиций токенов становится меньше, а записей словаря — больше. Именно такое соотношение показано на примере; он не доказывает оптимальность сжатия, уменьшение функции потерь предсказания следующего токена или большую осмысленность частей слов.

Будущая модель-декодер будет использовать конечный словарь ID для входов и выходов. BPE даёт этому словарю повторно используемые единицы содержимого. Обучение таблицы токенизатора отделено от обучения нейронных весов, с помощью которых модель будет описывать последовательности этих ID.

В главе 4 полученные правила будут зафиксированы и применены к новому входу без повторного обучения и подсчёта частот по валидационным или тестовым документам. Для BOS будет зарезервирован ID 0, для EOS — ID 1; каждый ID содержимого из главы 3 сдвинется на два. Поэтому ID 256 при обучении станет ID содержимого 258. Сохранённые байты токенов перейдут в следующую главу без изменений. Там будут разобраны применение зафиксированных правил и восстановление точных входных байтов.