Трансформер для структурированных документов Hierarchical Document Transformer
Механизм внимания в трансформерах дает большое преимущество — модель точно выбирает из контекста релевантные части. У этого механизма есть и довольно значимый недостаток — вычислительная сложность, которая квадратично возрастает в зависимости от размера входной последовательности. Одна из зацепок, которая может помочь облегчить работу вниманию — структура текста. Речь и про крупномасштабную структуру (вроде разделов, глав, параграфов), и про локальную структуру на уровне предложений.
Интуитивно понятно, что слова внутри одного предложения имеют друг к другу больше отношения, чем слова из далеких частей текста, хотя и далекие части осмысленного текста все же связаны. Это подтверждается и более строго — например, мы в свое время показали, что автокорреляция слов внутри одного текста убывает по степенному закону по мере отдаления друг от друга вне зависимости от языка. Причем неважно, “Дон Кихот” это или “Критика чистого разума”.
Может быть и такое, что мы, например, читаем книгу и вдруг понимаем, что этот момент встречался в какой-то другой главе. Но и тогда мы скорее всего обратимся сначала к содержанию, где найдем нужную главу, а затем уже в ней будем искать нужное предложение, а не просто листать назад всю книгу. То есть, даже обращаясь к далеким токенам, будем учитывать иерархическую структуру текста. Трансформеры же в основном воспринимают текст как единую плоскую массу.
Hierarchical Document Transformer (HDT) — новая модель, которая позволяет немного “расслабить” внимание трансформера, то есть проредить матрицу внимания с помощью информации об иерархической структуре.
Структура документа в HDT отражается с помощью вспомогательных якорных токенов. Например, документ просто разбит на параграфы. Тогда получим якорные токены трех видов — один будет соответствовать началу самого документа (желтый токен на картинке), второй — началу параграфов (красные токены), третий — началу предложений (синие токены). Сами слова — это обычные токены, которые оказываются каждый в своём блоке. Уровней может быть произвольное количество, авторы взяли три для простоты примера.
С помощью этой иерархии строится разреженная матрица внимания. Рядовые (зеленые) токены обмениваются информацией только со своими “собратьями”, то есть токенами своего блока, и “родительскими” токенами. Но это не означает, что токены из соседних параграфов вообще никак не связаны. Между ними тоже происходит обмен информацией, но не напрямую, а косвенно. Например, токены Т1 и Т3 связаны через токены SENT и SEC (то есть почти в буквальном смысле через “содержание”, так же, как и делал бы живой читатель длинного текста) Таким образом, информация из любого токена может перетечь в любой другой токен, но структура внимания гораздо более разреженная. Это даёт выигрыш с точки зрения вычислительной сложности и требуемой памяти. Теоретическая сложность вычислений с такой разреженной матрицей уменьшается c O(n^2) до О(n s), где s — длина самого длинного предложения в документа. Отсюда следует, что вычислительная выгода будет заметна для длинных документов, где s сильно меньше n. Однако, чтобы претворить эту пока теоретическую выгоду в жизнь, нужно подобрать правильный вид ядра внимания.
Чтобы достичь этого, авторы опирались на идею FlashAttenation. Там вычисления разделили на небольшие блоки токенов, размеров 128х64 и перевели их на статический RAM. Теперь, вместо того, чтобы просто делить матрицу на равные блоки, авторы усовершенствовали процесс с учетом структуры документа. Перед тем как скопировать K и V в память, их сортируют в соответствии с иерархией (документ-раздел-предложение-слово). В результате получается, что более связанные токены окажутся ближе друг к другу, а при дальнейшем перемножении на Q получится больше пустых блоков. Непустые блоки обрабатываются параллельно, а пустые просто пропускаются
HDT проверили на простых математических заданиях ListOps, на SciRepEval (оценка близости на научных статьях), FacetSum (суммаризация отдельных разделов научных статей) и SCROLLS (разнородные задания на длинных, не обязательно структурированных текстах). Причем так как генеральная цель работы — повысить эффективность вычислений, то обучали модель на скромных мощностях на корпусе текстов из unarXive, HUPD и википедии.
Обратим внимание на результаты на самом простом тесте — ListOps. Там стало видно, как влияет связь токенов разного уровня между собой (и как, следовательно, меняется матрица внимания). Описанный здесь пример связности соответствует типу r+g+b в таблице выше. Можно “отвязать” верхний уровень от следующих (g+b и blue). Это делает матрицу еще менее плотной и, как показывают эксперименты, увеличивает точность на этом датасете. На остальных датасетах HDT почти во всех случаях бьет текущий state-of-the-art, хотя и не слишком значительно.
Преимущества модели заметны на сравнении вычислительной эффективности. HDT запускали вместе с уже упомянутой FlashAttention и архитектурой Longformer. Стандартное внимание уже на 16 тысячах токенов съело более 40GB, поэтому выбыло из сравнения. Вычислительная выгода HDT всё более значима по мере увеличения числа токенов.
Можно утверждать, что отражение структуры документов в явном виде внутри механизма внимания действительно повышает эффективность. Важно также, что модель оказывается очень гибкой — можно менять уровни, их связность и добиваться разной плотности внимания. Наконец, возникает вопрос — что будет, если расширять иерархию вплоть до побайтового уровня? Можно ли прийти к token-free моделям?