Cell-NN

Краткое описание

Проект реализует экспериментальную архитектуру нейронной сети (НС) - НС с маршрутизацией на уровне отдельных узлов.

В проекте реализованы следующие архитектурные идеи НС:

• Блочная, а не слоистая структура сети: взаимодействия между элементами сети происходит на уровне блоков, а не слоев.
• Рекуррентность: блоки сети могут использоваться многократно в одном прямом проходе вычислений.
• Граф связей: связи между блоками задаются матрицей смежности, а не набором слоев.
• Маршрутизация встроена в каждый блок: отдельный блок "принимает решение" о том, куда передать выходной сигнал.
• Автоматическое формирование структуры НС по описанию (списку размерностей).

Использование

Скачать репозиторий, создать виртуальное окружение и установить зависимости, используя uv:

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

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

Created model with 8M parameters
Vocabulary size: 166
Vectorizing texts: 100%|████████████████████████████████████████████████████████████████| 1/1 [00:01<00:00,  1.09s/it]
Total 5943221 tokens from 5990428 chars
Packing into sliding window of size 1025
1: loss=3.9588 ppl=15.5495:   9%|█           | 131/1451 [00:16<02:50,  7.74it/s]

Используемые библиотеки и зависимости

Сознательным решением при создании проекта было использовать минимум сторонних библиотек python. Поэтому для реализации модели достаточно pytorch. Для визуализации и сериализации используются библиотеки matplotlib, networkx и dill.

В демонстрации языковых возможностей модели используются модули:

- tokenizelib.py - токенизация текста,
- datalib.py - работа с датасетами,
- trainer.py - обучение моделей.

Описание архитектурных решений

Основной pytorch-модуль модели - models.GridNet. Модель GridNet представляет собой экспериментальную нейронную сеть с блочной структурой и встроенной маршрутизацией сигналов на уровне отдельных узлов. Основными архитектурными решениями являются:

1. Блочная структура

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

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

2. Рекуррентность и повторное использование вычислительных блоков

Каждый блок (Cell + Router) может многократно использоваться в прямом проходе, благодаря чему модель эффективно обрабатывает данные с минимальным количеством параметров и улучшенной обобщающей способностью. Параметры блоков храняться в "общих" тензорах. Таким образом один тензор объединяет параметры нескольких блоков:

где

3. Графовая структура связей

Связи между блоками определяются в виде графа, представленного матрицей смежности специального вида. Узлы этого графа соответствуют блокам, а ребра задают возможные пути распространения сигнала. Такая структура позволяет модели динамически выбирать пути вычислений на основе входных данных.

Ниже пример матрицы смежности для регулярного графа из примера 1. Этот граф имеет 21 вершину, каждая из которых связана с двумя другими, кроме вершин последнего слоя. Поэтому матрица смежности состоит из 21-й строки, каждая соответствует вершине графа. Номер столбца соответствует номеру связи. На пересечении i-й строки и j-го столбца - номер вершины, с которой связана i-я вершина j-й дугой.

tensor([[ 2,  3],
        [ 4,  5],
        [ 6,  7],
        [ 8,  9],
        [10, 11],
        [12, 13],
        [14, 15],
        [16, 17],
        [18, 19],
        [20, 21],
        [14, 15],
        [16, 17],
        [18, 19],
        [20, 21],
        [-1, -1],
        [-1, -1],
        [-1, -1],
        [-1, -1],
        [-1, -1],
        [-1, -1],
        [-1, -1],
        [-1, -1]])

4. Встроенная маршрутизация

Каждый каждый блок сети содержит вычислитель Сell и маршрутизатор Router. Вычислитель осуществляет прямое преобразование, а маршрутизатор решает, по какому пути направить сигнал, используя обучаемые веса и оценки значимости путей. Это решение принимается на основе входного сигнала и текущего состояния графа. Подобный подход аналогичен схеме Mixture-of-Experts, однако в отличие от последнего работает не на уровне маршрутизации выхода слоя, а на уровне каждого отдельного блока сети. Подход обеспечивает модели возможность адаптивной обработки данных, улучшая её производительность в задачах с неоднородными или сложными входными структурами.

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

5. Автоматическое формирование структуры сети

Архитектура сети автоматически формируется по минимальному описанию, включающему только размеры ячеек, степень связности графа и количество уровней решётки (lattice). Это позволяет легко настраивать и масштабировать модель без необходимости ручного проектирования сложных структур.

Пример формирования структуры из списка уровней для графа из примера выше.

Преимущества архитектуры GridNet

1. Линейный, а не квадратичный рост необходимой памяти при увеличении количества вычислительных блоков. Каждый новый блок требует хранения информации о связях с k соседями, а не со всеми элементами следующего слоя (как, например в FFN), и не квадратичный рост от роста размера матриц Q, K, V в Attention блоке.
2. Принцип локальной маршрутизации вычислений, использованный в Mixture-of-Experts, исполнен на более глубоком уровне. Вместо выбора нескольких экспертов на каждом слое, выбор осуществляется в каждом блоке и адресуемый блок может быть не только в следующем слое, но и любым блоком модели.
3. Связи между блоками задаются графом. Связи между блоками задаются в виде матрицы смежности размерности
   I x K
   : в k-й колонке i-й строки записан индексы блока смежного с i-м.
4. Возможность переиспользование блоков за счет рекуррентных связей. Структура графа вычислений не ограничена слоистыми архитектурами, а может содержать циклы.
5. Нет позиционного кодирования в задачах обработки текста. Элементы входного тензора (эмбеддинги токенов) подаются на входные блоки сети последовательно, в соответствии с естественным порядком индексов. Таким образом позиции токенов выражаются топологически, - номер блока обработчика элемента соответствует его позиции в последовательности. Это отличется от принятого в трансформера кодирование позиции токена через изменение значений эмбеддинга этого токена (как например, в RoPE). Это позволяет уменьшить количество вычислений и сделать ычисления более прозрачными.