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