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

Сначала рассчитайте два раунда вручную

Даны два отдельных документа. Не соединяйте их в одну строку:

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

До первого слияния каждый токен представляет один исходный байт, а его ID равен десятичному значению этого байта. Поэтому ID 97 соответствует байту 61 в шестнадцатеричной записи и букве a; ID 98 соответствует байту 62 и букве b.

Перечислите все позиции, с которых начинаются соседние пары. Нумерация начинается с нуля. В документе train-aaa пара (97,97) начинается дважды: в позициях 0 и 1. Эти два вхождения используют общий средний токен 97, то есть перекрываются, но при выборе правила учитываются оба. Документ train-aba добавляет одно вхождение (97,98) и одно вхождение (98,97):

Первое обученное правило получает ранг 0. На этом раунде чаще всего встречается пара (97,97), поэтому ей назначается новый ID 256.

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

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

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

Теперь выполните расчёт для ранга 1. После первой замены пары (97,98), (98,97) и (256,97) встречаются по одному разу. При равной частоте курс сравнивает пары лексикографически по числовым ID: сначала левые ID, затем правые. У (97,98) левый ID равен 97, поэтому эта пара опережает пары с левыми ID 98 и 256. Новому токену 257 соответствует последовательность байтов 61 62. После замены получаем:

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

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

Как однозначно выбрать следующее правило

На каждом раунде выбранную пару и соответствующий ей новый токен определяет формула:

Для каждой пары-кандидата строится тройка $(C(a,b),-a,-b)$, а затем такие тройки сравниваются лексикографически. Сначала выбирается наибольшая частота $C(a,b)$. Если частоты равны, наибольшее значение $-a$ соответствует меньшему левому ID. Если совпадают и частота, и левый ID, наибольшее значение $-b$ соответствует меньшему правому ID. Это правило обеспечивает воспроизводимость именно в этом курсе; другие реализации BPE могут выбирать победителя при равной частоте по другому правилу.

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

Если в начале очередного раунда ни в одном документе нет соседней пары, обучение завершается раньше запрошенного числа раундов. Частота 1 сама по себе не служит причиной для остановки: такое правило можно обучить.

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

До начала слияний каждому возможному значению байта уже соответствует свой ID из диапазона 0..=255, то есть словарь токенов изначально содержит 256 исходных токенов. Успешный раунд с номером $r$ при нумерации с нуля создаёт токен с ID $256+r$. Поэтому после $k$ успешных раундов словарь токенов главы 3 содержит $256+k$ токенов. В главе 4 эти ID содержимого будут перенесены в окончательную схему после резервирования управляющих токенов.

Почему словаря целых слов оказалось недостаточно

Фиксированный словарь назначает отдельный ID каждому слову, которое встретилось при его построении. Для любого незнакомого написания нужен запасной ID. В демонстрационной программе ниже слово lower есть в словаре и получает собственный ID, а lowering и все остальные незнакомые слова получают общий ID 0. Модель уже не может отличить одно такое слово от другого:

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 году Филип Гейдж предложил для сжатия данных другой алгоритм. Он находил частую соседнюю пару байтов, заменял её одним ещё не использованным значением байта, а сведения о замене сохранял для последующего восстановления данных. Подробности приведены в исходной статье о сжатии парами байтов. Это не токенизатор из курса: алгоритм Гейджа ограничен свободными значениями байтов, тогда как наша реализация использует ID типа u32 и хранит точную последовательность байтов для каждого нового токена.

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

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

Главная мысль: обучение BPE состоит из последовательности точных решений: подсчитать пары-кандидаты, выбрать победителя при равной частоте и выполнить замену без перекрытий. Фраза «сливать частые пары» скрывает эти необходимые шаги.

Реализуйте обучение BPE без готового токенизатора

Счётчик получает срез последовательностей токенов — по одной последовательности на документ, — а не один общий массив. Метод windows(2) вызывается отдельно для каждой последовательности. Поэтому он учитывает оба перекрывающихся окна в aaa, но по самой структуре данных не может построить окно через границу двух документов.

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

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 и увеличивает индекс на 2. Если совпадения нет, она копирует текущий ID и увеличивает индекс на 1:

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, а не произвольный список документов. Вызов training_documents() возвращает только обучающую выборку. Метод сохраняет стабильные ID именно этих документов и преобразует в начальные 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)
    }

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

Функция main из демонстрационной программы печатает первые восемь правил, полученных при обучении на корпусе, в порядке их рангов, а затем короткую трассировку двух раундов для 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(())
}

Запустите команды из корня репозитория:

cargo test --workspace --locked
cargo run --quiet --locked -p ch03-learn-bpe-merges | diff -u rust/demos/ch03-learn-bpe-merges/expected.txt -

Если вывод совпадает с ожидаемым, diff ничего не напечатает. На реальной обучающей выборке правило ранга 0 объединяет пару (32,208). Она встречается 81 раз, получает новый ID 256 и хранит байты 20 d0. Байт 20 — это пробел. d0 — первый байт многих двухбайтовых последовательностей UTF-8 для кириллических символов. Здесь после него нет байта продолжения, поэтому последовательность 20 d0 остаётся незавершённой и не является ни отдельным символом, ни корректной самостоятельной строкой UTF-8. Показывать эту последовательность нужно как байты или шестнадцатеричные значения.

Практический вывод: сохраняйте точную последовательность байтов вместе с каждым новым ID. Тогда декодирование сможет заменить ID именно теми байтами, из которых был построен токен.

Сверьте ручной расчёт со схемой

На схеме показана короткая трассировка для aaa и aba, а не восемь правил, полученных при обучении на реальном корпусе. Каждый этап содержит оба документа. В таблице очередного раунда число вхождений пары с учётом перекрытий отделено от числа замен без перекрытий. Выбранная строка помечена текстом «Пара выбрана», звёздочкой, жирным начертанием и рамкой, поэтому результат можно определить, не опираясь только на цвет. Числовые ID и шестнадцатеричные байты читаются слева направо при любом языке страницы; на узком экране этапы располагаются друг под другом.

Вернитесь к этапу 0 и найдите средний токен 97 в документе train-aaa. Он входит сразу в два перекрывающихся вхождения (97,97). Поэтому на ранге 0 число вхождений равно 2, а число замен — 1. Для ранга 1 не используйте положение строк в таблице как подсказку: сравните числовые пары по правилу из формулы. Пара $(97,98)$ должна опередить $(98,97)$ и $(256,97)$.

Сначала решите задачи, затем сверьте ответы

1. Для одного документа [97,97,97,97] найдите $C(97,97)$ и число замен за один проход слева направо.
2. После раунда 0 в учебном примере запишите тройку оценки для каждого из трёх кандидатов с частотой 1 и выберите победителя.
3. Объясните, почему два документа [97] | [98] не создают пару-кандидат (97,98).
4. Найдите размер словаря токенов, если было запрошено 12 раундов, но обучение остановилось после 9 успешных.
5. Определите, изменят ли 10 000 дополнительных повторов в валидационном документе таблицу правил, полученную при обучении.
6. Запишите байты результата гипотетического слияния (256,97), если токен 256 хранит байты 61 61.
7. Объясните, почему байты реального правила ранга 0 20 d0 нельзя вывести как один символ.

Если вы меняете реализацию, сначала предскажите первый ранг, который должен измениться. При несовпадении с ожидаемым выводом проверяйте причины по отдельности: какие документы вошли в обучение, где проходят их границы, как считаются вхождения, как выбирается победитель при равных частотах и как выполняется замена. Одного слова «BPE» недостаточно, чтобы считать разные реализации эквивалентными.

Зафиксируйте правила перед кодированием

В этой главе мы только обучили и упорядочили правила слияния. Способ кодирования произвольных входных данных ещё не определён. Обучение задаёт порядок правил; во время последующего кодирования уже зафиксированные правила применяются к новому тексту без повторного подсчёта частот.

В главе 4 для BOS и EOS будут зарезервированы два ID, поэтому каждый ID содержимого из главы 3 сдвинется на два. Затем правила будут применяться к произвольным байтам UTF-8 по возрастанию ранга: правило с меньшим рангом может создать токен, который понадобится правилу с большим рангом. При декодировании каждый ID будет заменён сохранённой для него последовательностью байтов, что восстановит исходные байты без изменений. Валидационные и тестовые документы можно использовать для проверки обратного преобразования, но их содержимое не участвует в подсчёте пар и обучении правил.