Meta DLRM

01 · CONTEXTO

Um modelo de recomendação que não cabe em um dispositivo

Modelos de recomendação baseados em redes neurais diferem de outras redes profundas porque precisam trabalhar com dados categóricos. O blog de IA da empresa conta que a falta de detalhes públicos de modelos e de conjuntos de dados representativos atrasava a comunidade de pesquisa, e anuncia a abertura do código do DLRM. O artigo traz a implementação em PyTorch e em Caffe2.

Esses sistemas são usados em tarefas como a previsão de taxa de cliques em anúncios e o ranqueamento.

02 · ARQUITETURA

Do dado categórico à probabilidade de clique

O DLRM trata dois tipos de entrada. As características categóricas passam por embeddings; as contínuas, por uma MLP inferior. Depois as interações de segunda ordem são calculadas de forma explícita, e o resultado passa por uma MLP superior e por uma sigmoide, que dá a probabilidade de clique.

A interação é o produto escalar entre todos os pares de vetores de embedding e das características densas já processadas. Segundo o artigo, isso imita máquinas de fatoração e reduz de forma significativa a dimensionalidade do modelo, por considerar só termos cruzados.

1. As características categóricas são processadas por embeddings.

2. As características contínuas passam pela MLP inferior.

3. A camada de interação calcula o produto escalar entre pares de vetores de embedding e das características densas processadas.

4. A MLP superior processa as interações e uma sigmoide devolve a probabilidade de clique.

03 · DOIS GARGALOS

Embeddings pesam em memória, MLPs pesam em cálculo

Segundo o artigo, o DLRM tem até várias ordens de magnitude mais parâmetros do que outros modelos comuns, como as redes convolucionais, e o treino pode levar semanas ou mais. Por isso é importante paralelizá-lo de forma eficiente.

04 · PARALELISMO COMBINADO

Repartir as embeddings, replicar as MLPs

O esquema do artigo combina os dois: paralelismo de modelo nas tabelas de embedding, para aliviar a restrição de memória, e paralelismo de dados nas camadas totalmente conectadas, para escalar o cálculo. Segundo o artigo, essa combinação é um requisito único do DLRM, e o artigo diz que nem o Caffe2 nem o PyTorch suportam essa combinação.

1. As tabelas de embedding ficam repartidas entre os dispositivos, e cada um guarda só uma parte delas.

2. Ao fim da busca, cada dispositivo tem os vetores das tabelas que guarda, mas a MLP superior e a interação precisam de todos os embeddings.

3. A comunicação all-to-all, chamada de butterfly shuffle, redistribui o resultado da busca do minibatch inteiro em partes do minibatch para todos os dispositivos.

4. Cada dispositivo processa amostras diferentes com uma réplica das MLPs, que é o paralelismo de dados.

5. Na volta, as atualizações de parâmetros das MLPs são acumuladas com allreduce e aplicadas às réplicas de modo síncrono, deixando os parâmetros consistentes antes de cada iteração.

6. Na versão descrita no artigo, as transferências do shuffle são cópias explícitas, e os autores pretendem otimizá-las.

05 · MEDIÇÃO E LIMITES

Como o benchmark é usado e onde ele para

O artigo apresenta o DLRM como benchmark para experimentação algorítmica e co-design de sistemas. O benchmark gera entradas aleatórias e sintéticas; as sintéticas permitem, por exemplo, não compartilhar um conjunto de dados por privacidade, expressando as características categóricas por distribuições.

Com dados públicos, o conjunto Kaggle de anúncios tem 13 características contínuas e 26 categóricas, o que define o tamanho da entrada da MLP e o número de embeddings do modelo.

No perfil de um modelo de exemplo, a maior parte do tempo vai para a busca de embeddings e para as camadas totalmente conectadas. Na CPU essas camadas pesam bastante, e na GPU são quase desprezíveis.