DBRX (modelo de linguagem)

DBRX é um grande modelo de linguagem (LLM) aberto, desenvolvido pela equipe de pesquisa Mosaic AI da empresa Databricks. O modelo foi lançado oficialmente em 27 de março de 2024 e é posicionado como uma solução de alto desempenho para uso corporativo.

O DBRX é construído sobre uma arquitetura de mistura de especialistas (MoE) de granularidade fina e combina alto desempenho com eficiência de treinamento e inferência. No momento de seu lançamento, o DBRX demonstrou os melhores resultados entre todos os modelos abertos nos principais benchmarks, superando modelos como LLaMA 2, Mixtral e Grok-1, e mostrando-se competitivo com modelos fechados do nível do GPT-3.5 Turbo.

Histórico de Desenvolvimento

O surgimento do DBRX foi uma continuação da estratégia da Databricks de desenvolver modelos generativos abertos. Em junho de 2023, a Databricks adquiriu a startup MosaicML, especializada no treinamento de grandes modelos, e, com base nela, foi criada a divisão Mosaic AI.

A equipe da Mosaic AI, liderada pelo arquiteto-chefe de redes neurais Jonathan Frankle, iniciou o desenvolvimento de um novo grande LLM com o objetivo de alcançar uma qualidade comparável aos melhores sistemas proprietários, mas em formato aberto. O projeto foi nomeado DBRX. O desenvolvimento e o pré-treinamento do modelo levaram cerca de 2,5 meses e, segundo estimativas, custaram aproximadamente $10 milhões.

Arquitetura

O DBRX é um modelo transformer do tipo somente decodificador (decoder-only) e implementa uma arquitetura de mistura de especialistas (MoE) de granularidade fina (fine-grained).

Principais características da arquitetura:

• Número total de parâmetros: 132 bilhões.
• Especialistas: O modelo é composto por 16 pequenos submodelos especializados ("especialistas").
• Mecanismo de ativação: Para cada token de entrada, apenas 4 dos 16 especialistas são ativados. Isso significa que, durante a inferência, apenas 36 bilhões de parâmetros estão ativos, o que garante alta velocidade e eficiência. Esse esquema oferece 65 vezes mais combinações possíveis de especialistas em comparação com o modelo Mixtral (8 especialistas com ativação de 2).
• Componentes: São utilizadas soluções arquitetônicas modernas, como embeddings de posição rotacional (RoPE), gated linear units (GLU) e grouped query attention (GQA).
• Comprimento do contexto: 32.768 tokens.

Essa arquitetura permite que o modelo combine as vantagens de um número enorme de parâmetros (para armazenamento de conhecimento) com a eficiência de modelos menores (para velocidade de inferência).

Treinamento

O pré-treinamento do DBRX foi realizado em um conjunto de dados cuidadosamente curado de 12 trilhões de tokens, composto por textos e código. A qualidade dos dados foi uma prioridade fundamental: os desenvolvedores utilizaram a plataforma em nuvem da Databricks (Apache Spark, Databricks Notebooks, Unity Catalog) para limpar, preparar e auditar os dados.

Durante o treinamento, foi aplicado o método de aprendizagem curricular (curriculum learning), no qual a proporção dos tipos de dados foi alterada em diferentes estágios. Por exemplo, a parte final do treinamento foi dedicada à introdução gradual de tarefas complexas, o que, segundo os desenvolvedores, resultou em um aumento significativo na qualidade. O treinamento foi realizado em um cluster de 3.072 GPUs Nvidia H100.

Após o pré-treinamento, o modelo base passou por um ajuste fino adicional (instruction tuning) para criar a versão interativa DBRX Instruct, otimizada para seguir as instruções do usuário.

Desempenho

No momento de seu lançamento, o DBRX estabeleceu um novo padrão de qualidade para LLMs abertos em uma ampla gama de benchmarks.

Comparação com modelos abertos

Resultados do DBRX Instruct nos principais benchmarks

Benchmark	Tarefa	DBRX Instruct	Próximo melhor (Mixtral/Grok-1)
Hugging Face Open LLM Leaderboard (AVG)	Conhecimento geral	74,5%	72,7% (Mixtral Instruct)
HumanEval	Programação	70,1%	63,2% (Grok-1)
GSM8K	Raciocínio matemático	66,9%	62,9% (Grok-1)
MMLU	Conhecimento geral	73,7%	71,5% (Mixtral Instruct)

O DBRX conquistou o primeiro lugar tanto no ranking geral do Hugging Face Open LLM Leaderboard quanto no teste abrangente Databricks LLM Gauntlet, demonstrando uma vantagem significativa sobre seus predecessores.

Comparação com modelos fechados

O DBRX Instruct supera o GPT-3.5 Turbo em vários indicadores-chave, incluindo MMLU (73,7% contra 70,0%) e HumanEval (70,1% contra 48,1%). Em termos de qualidade de resposta em alguns benchmarks (por exemplo, MTBench), o modelo se aproxima do nível do Gemini 1.0 Pro e das primeiras versões do GPT-4.

Eficiência de Treinamento e Inferência

• Eficiência de treinamento: O uso da arquitetura MoE permitiu reduzir os custos em FLOPS de 2 a 4 vezes em comparação com modelos densos de qualidade similar.
• Eficiência de inferência: Ao ativar apenas 36 bilhões de parâmetros, o DBRX oferece uma taxa de transferência (velocidade de inferência) 2 a 3 vezes maior em comparação com modelos densos de tamanho equivalente (por exemplo, LLaMA2-70B).

Licenciamento e Disponibilidade

O DBRX é distribuído sob uma licença especialmente desenvolvida, a Databricks Open Model License. Ela permite o uso e a modificação livres, incluindo para fins comerciais, mas contém algumas restrições. Em particular, assim como a licença do LLaMA 2, ela exige uma permissão separada da Databricks se os serviços baseados no DBRX forem utilizados por uma audiência superior a 700 milhões de usuários ativos mensais.

Os pesos pré-treinados do modelo (versões base e Instruct) estão disponíveis para download no repositório do Hugging Face.

Literatura

• Mosaic Research Team. (2024). Introducing DBRX: A New State‑of‑the‑Art Open LLM. Databricks Blog.
• Databricks. (2024). Databricks Open Model License (DBRX). Online specification.
• Fedus, W.; et al. (2021). Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. arXiv:2101.03961.
• Lepikhin, D.; et al. (2020). GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding. arXiv:2006.16668.
• Ainslie, J.; et al. (2023). Grouped‑Query Attention: Efficient Training of Generalized Multi‑Query Transformers. arXiv:2305.13245.
• Su, J.; et al. (2021). RoFormer: Enhanced Transformer with Rotary Position Embedding. arXiv:2104.09864.
• Dao, T. (2023). FlashAttention‑2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691.
• Cai, W.; et al. (2024). A Survey on Mixture of Experts in Large Language Models. arXiv:2407.06204.
• Chen, Y.; et al. (2024). Scaling Laws for Fine‑Grained Mixture of Experts. arXiv:2402.07871.
• Kundu, A.; et al. (2024). Strategic Data Ordering: Enhancing Large Language Model Training via Curriculum Learning. arXiv:2405.07490.