🌟 MoBA: Метод эффективной обработки длинных контекстов в LLM.

Mixture of Block Attention (MoBA) - метод, разработанный MoonshotAI для повышения эффективности обработки длинных последовательностей в LLM. MoBA основан на принципах Mixture of Experts и применяется к механизму внимания в архитектуре Transformers. Он позволяет динамически выбирать исторически релевантные KV-блоки для каждого токена запроса, снижая, как следствие, вычислительные затраты при обработке длинных контекстов.

MoBA разделяет контекст на блоки и использует механизм маршрутизации для выбора наиболее релевантных блоков. Такая конструкция помогает адаптивно фокусироваться на информативных частях контекста, что полезно для задач, требующих обработки длинных документов. Метод сохраняет причинность (causality) в авторегрессионных моделях за счет ограничения внимания только текущими и прошлыми блоками.

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

Эксперименты показали, что MoBA имеет сопоставимую производительность с Full attention при значительно меньших вычислительных затратах. Например, на Llama-8B-1M-MoBA с длиной контекста до 1 млн. токенов MoBA достигает разреженности до 95.31%, при этом сохраняя высокую точность на бенчмарках (AGIEval, BBH, CEval и др.).

На тестах с RULER с длиной контекста 128K MoBA показал результат 0.7818, что близко к результату полного внимания (0.7849).

⚠️ Актуальная реализация ядра полагается на flash-attn= =2.6.3. Данная реализация MoBA полностью совместима с transformers. Выбор бекэнда выполняется параметрами --attn moba и --attn moba_naive

▶️Локальная установка и запуск:

# Clone the repository
git clone https://github.com/MoonshotAI/MoBA.git

# Create a Conda venv
conda create -n moba python=3.10
conda activate moba

# Install dependencies
pip install .

# Quick Start 
python3 examples/llama.py --model meta-llama/Llama-3.1-8B --attn moba

# Unit Tests
pytest tests/test_moba_attn.py

📌Лицензирование: MIT License.


🟡Техотчет
🖥GitHub


@ai_machinelearning_big_data

#AI #ML #LLM #MoBA