Token-estimation-function (TEF)

Проект исследует и реализует модуль функции оценки токенов (token estimation function) на основе модели GPT-2.

Идея TEF

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

Другая интерпретация: каждая оценка токена представляет его объясненную дисперсию. Фильтрация токенов происходит по порогу объясненной дисперсии.

Related Works

Как и в Mixture of Experts (MoE), модель должна активировать только вычислительно важные части сети, но в случае TEF роутинг выполняется на уровне токенов через TEF-шлюзы (или через фильтры на основе кумулятивной объясненной дисперсии).

Fast Feedforward Networks

papers/Fast Feedforward Networks.pdf

Работа фокусируется на ускорении наиболее дорогого компонента Transformer-блоков - FFN части. FFN блоки без MoE архитектуры доминируют по количеству FLOPs над остальными блокам. Таким образом, токен-ориентированный роутинг на уровне целого трансфомерного слоя становится естественным способом повысить разреженность и снизить итоговую вычислительную нагрузку без построения MoE.

HyperAttention: Long-context Attention in Near-Linear Time

papers/HyperAttention - Long-context Attention in Near-Linear Time.pdf

HyperAttention показывает, что внимание можно считать существенно экономнее за счет алгоритмической структуризации и разрежения взаимодействий. Эти идеи близки к TEF: TEF не меняет формулу внимания напрямую, но вводит обучаемую разреженность по токенам, уменьшая эффективный объем вычислений в блоках и приближая поведение к sparse-attention режиму.

Infini-attention: Infinite Context with Bounded Memory

papers/Infini-attention - Infinite Context with Bounded Memory.pdf

Infini-attention развивает идею селективного сохранения информации и ограниченного бюджета вычислений/памяти при длинном контексте. Это похоже на TEF - не все токены должны быть одинаково активны на каждом шаге, и модель выигрывает от механизма, который адаптивно определяет важность и оставляет только информативную часть сигнала.

FTP: A Fine-Graned Token-Wise Pruner For Large Language Models via Token Routing

Наиболее близкая к TEF по идеям работа. Однако, вместо функции оценки, используется роутер. Причем, если в GPT2-TEF перед каждым слоем трансфомера свой уникальный модуль функции оценки, в FTP используется единый роутер для всех слоёв.

Важно, что в работе также доказывается, что после примерно третьего или четвертого слоя, важность обработки каждого токена последовательности значительно снижается, поскольку изменения токена после обработки минимальны.

Связь с текущей работой

В репозитории акцент идет на разреженности на уровне токена, а не на уровнях кеша, FFN или внимания. TEF действует скорее как роутер (аналог MoE-гейтинга, но для токенов); Поддерживаются разные режимы разрежения (off / eval / train) и используется дополнительная регуляризация разреженности (GatedL1Loss), чтобы управлять компромиссом между качеством и экономией вычислений.

Итого, выбранные работы задают теоретический и практический контекст для условных вычислений и разреженных преобразований в Transformer, а TEF в этом проекте можно рассматривать как шаг к унифицированному sparsity-механизму на уровне токенов.

Реализация

1. Интерфейс модуля функции оценки токенов: src/tefs/tef_scorer.py
2. Бэкенд: src/tefs_backend/tef_scorer/
3. Модель GPT-2 со встроенной функцией оценки токенов: src/models/gpt2_tef.py

Бэкенд

Реализация бэкнда на torch src/tefs_backend/tef_scorer/torch_backend.py и на основе triton src/tefs_backend/tef_scorer/triton_backend.py

На triton реализован кернел, "раскидывающий" токены после сортировки по нужным позициям.

Triton-реализация дает ~5% ускорения одной итерации на NVIDIA 3090

Датасет

WikiText-103 — это датасет для языкового моделирования из статей Википедии на английском языке.

Он содержит 103 млн токенов в train, составляющие длинные связные текста, а не отдельные несвязные короткие предложения.

Структура проекта

Основные директории:

• src/models/ - обучаемые модели.
• src/lightning_modules/ - Lightning-модули и обертки над моделями.
• src/data/ - директория для датасетов (представлен только WikiText-103).
• src/tefs/ и src/tefs_backend/ - реализация TEF.
• src/losses/ - aux loss'ы для обучения TEF.
• configs/ - YAML-конфиги моделей, обучения и генерации.
• scripts/training/ - скрипты для запуска обучения. Все частные скрипы опираются на общий скрипт train.py
• scripts/evaluation/ - скрипты для генерация текстов и оценки результатов обучения.
• scripts/benchmarks/ - бенчмарки для оценки triton-кернелов.
• scripts/profiling/ - скрипт для профилирования.
• tests/ - unit-тесты на pytest.

Установка

Все зависимости для минимально работоспособной версии находятся в файле requirements.txt. Дополнительные зависимости (для разработки и тестирования) лежат в requirements_dev.txt

Для обучения на видеокартах NVIDIA нужно установить версию torch под свою конфигурацию.

Старт

Для скриптов обязательно нужно задать .env с CONFIG_DIR:

CONFIGS_DIR="<путь до корня проекта>/configs"

Обучение baseline

python scripts/training/autoregressive/train_gpt2.py

Обучение GPT-2 + TEF

python scripts/training/autoregressive/train_gpt2_tef.py

Логирование и артефакты

Логи и артефакты:

• TensorBoard: logs/<task_name>/...
• Checkpoints: checkpoints/<project_name>/<task_name>/...

Также в скриптах подключен ClearML

Автотесты

Реализованы юнит-тесты на pytest. Из корня проекта:

pytest tests

Постер

Репозиторий - часть проекта на AITH DemoDay. Постер с демо-дей лежит в papers/