TikTok/ByteDance Monolith (+ For You)

01 · CONTEXTO

Recomendação que precisa aprender em tempo real

O Monolith é um sistema de recomendação descrito em um paper de autores ligados à ByteDance. O paper abre dizendo que um sistema de recomendação escalável e em tempo real é vital para muitos negócios movidos por feedback do cliente sensível ao tempo, como o ranqueamento de vídeos curtos e os anúncios online.

O ponto de partida do paper é que os frameworks de aprendizado profundo de uso geral ficam aquém desse cenário. Um dos motivos é que eles separam por completo o estágio de treino em lote e o de serving, o que impede o modelo de interagir com o feedback do cliente em tempo real.

Os dados de recomendação diferem de problemas como linguagem ou visão computacional em dois aspectos: as features são majoritariamente esparsas, categóricas e mudam de forma dinâmica, e a distribuição dos dados de treino não é estacionária, fenômeno que o paper chama de concept drift.

Como a mesma pessoa pode mudar de interesse a cada minuto, o paper defende atualizar os modelos em serving com o feedback novo o mais perto possível do tempo real.

02 · ARQUITETURA

Worker, servidor de parâmetros e dois estágios de treino

A arquitetura segue o modelo distribuído Worker-ParameterServer do TensorFlow: máquinas worker executam as computações do grafo, e máquinas PS guardam os parâmetros e os atualizam conforme os gradientes. Os parâmetros se dividem em dois conjuntos: os densos, que são os pesos da rede neural, e os esparsos, que são as tabelas de embedding das features esparsas.

O treino é dividido em dois estágios. O estágio em lote retreina com dados históricos, útil quando a arquitetura do modelo muda. Depois do deploy para serving, o treino não para: no estágio online, o worker consome dados em tempo real e atualiza o training PS, que sincroniza seus parâmetros com o serving PS periodicamente.

1. No estágio em lote, o worker treina com dados históricos; isso serve para retreinar quando a arquitetura do modelo muda.

2. Depois que o modelo entra em serving, o treino não para: no estágio online, o worker passa a consumir dados em tempo real.

3. O worker atualiza o training PS com o que aprendeu dos dados novos.

4. O training PS sincroniza periodicamente seus parâmetros com o serving PS.

5. O servidor de modelo responde ao usuário com os parâmetros do serving PS, e a sincronização passa a valer imediatamente do lado do usuário.

03 · EMBEDDINGS

Uma tabela de embeddings sem colisão

A tabela de embeddings de um sistema de recomendação é enorme: usuários e itens são ordens de grandeza mais numerosos que as palavras de um modelo de linguagem, e a tabela deve crescer conforme novos IDs são admitidos. Uma tabela assim dificilmente caberia na memória de um único host, e alguns frameworks a representam com variáveis densas de tamanho fixo.

Muitos sistemas contornam isso com hashing, que parte da suposição de que os IDs têm frequência uniforme e de que colisões são inofensivas. Segundo o paper, com o crescimento da tabela a chance de colisão de chaves aumenta e a qualidade do modelo se deteriora.

O Monolith desenvolveu uma HashTable chave-valor para os parâmetros esparsos, construída sobre Cuckoo Hashmap, que insere chaves novas sem colidir com as existentes. Consultas e atualizações da tabela também são operações nativas do TensorFlow.

Inserir todo ID novo esgotaria a memória rapidamente, então a tabela filtra. O paper observa que IDs raros contribuem pouco para a qualidade e que IDs antigos raramente contribuem para o modelo atual. Há dois métodos de filtragem antes da admissão: um limiar de ocorrências, ajustável por modelo, e um filtro probabilístico. Além disso, os IDs expiram depois de um período de inatividade, ajustável por tabela.

04 · MOTOR DE STREAMING

Kafka, Flink e o online joiner

O motor de streaming usa uma fila Kafka para registrar as ações dos usuários e outra para as features. No centro está um job Flink, o online joiner, que concatena as features com os rótulos das ações e produz exemplos de treino, escritos em mais uma fila Kafka. Essa fila é consumida tanto pelo treino online quanto pelo treino em lote, e o desenho permite alternar entre os dois.

1. As ações dos usuários vão para uma fila Kafka de ações.

2. As features seguem para outra fila Kafka, a de features.

3. O joiner, um job Flink, concatena as features com os rótulos das ações e produz exemplos de treino.

4. Ações e features chegam sem garantia de ordem temporal, e uma chave única por requisição permite parear as duas.

5. Quando a ação chega, o joiner consulta primeiro o cache em memória e depois o armazenamento chave-valor em disco, onde as features esperam.

6. Os exemplos vão para a fila de exemplos: o treino online lê direto dela, e o treino em lote passa pelo HDFS.

O joiner enfrenta problemas práticos. Ações e features chegam sem garantia de ordem temporal, então cada requisição leva uma chave única que permite parear as duas. E uma ação pode chegar alguns dias depois, enquanto manter todas as features em cache simplesmente não caberia na memória. Por isso há um armazenamento chave-valor em disco: quando um log de ação chega, o joiner consulta primeiro o cache em memória e, se não encontrar, o armazenamento em disco.

Há ainda o desbalanceamento entre classes: os exemplos negativos podem ser ordens de magnitude mais numerosos que os positivos. Uma estratégia comum é a amostragem negativa, que muda a distribuição do modelo treinado; para compensar, o Monolith aplica uma correção de log odds no serving, de modo que o modelo online seja um estimador não enviesado da distribuição original.

05 · SINCRONIZAÇÃO E FALHAS

Sincronizar sem parar o serving e aceitar perdas

Os modelos em produção costumam ter vários terabytes, e o serving PS não pode parar de servir enquanto é atualizado. Transferir o modelo inteiro do training PS pressionaria a rede e a memória do PS, porque receber o modelo novo exigiria o dobro do tamanho do modelo em memória.

A solução do Monolith é uma sincronização incremental. Um conjunto de chaves tocadas guarda os IDs cujos embeddings foram treinados desde a última sincronização, e só esse subconjunto dos parâmetros esparsos vai do training PS ao serving PS, em intervalos de minutos. O pacote incremental é leve para a rede e não causa pico de memória.

1. Só uma parte dos embeddings é treinada por vez, e o training PS guarda as chaves tocadas desde a última sincronização.

2. Em intervalos de minutos, só os embeddings tocados seguem para o serving PS; o pacote é leve para a rede e não causa pico de memória.

3. O serving PS continua atendendo o servidor de modelo enquanto recebe as atualizações.

4. Os parâmetros densos mudam mais devagar e são sincronizados com menos frequência, o que pode deixá-los numa versão mais defasada que a esparsa.

5. Se um training PS falha, ele perde um dia de atualizações, porque o snapshot de todos os training PS é feito todo dia.

Os parâmetros densos mudam mais devagar que os embeddings esparsos, então são sincronizados com menos frequência, em nível de dia. Isso pode deixar a parte densa numa versão mais defasada que a esparsa; o paper afirma que essa inconsistência pode ser tolerada.

Para falhas, o Monolith faz snapshot de todos os training PS todos os dias. Em caso de falha, um PS perde um dia de atualizações, mas a degradação de desempenho foi tolerável nos experimentos do paper. O paper apresenta isso como uma troca de confiabilidade do sistema por aprendizado em tempo real.

Nos experimentos de sincronização, a qualidade do modelo melhora de forma consistente à medida que a sincronização de parâmetros aumenta, e a conclusão do paper é que o intervalo deve ser o mais curto possível. O próprio paper limita isso ao que o custo de computação e a confiabilidade aguentam.

06 · O FEED FOR YOU

O que a página pública do TikTok descreve

A página da newsroom do TikTok, publicada em junho de 2020, explica que o feed For You é alimentado por um sistema de recomendação que entrega a cada pessoa o conteúdo com provável interesse. Ela trata de sinais e regras do produto e não descreve a infraestrutura de treino.

Entre os fatores estão as interações do usuário, como vídeos curtidos ou compartilhados, contas seguidas, comentários e conteúdo criado; as informações do vídeo, como legendas, sons e hashtags; e as configurações de dispositivo e conta, que recebem peso menor que os demais pontos medidos.

Um indicador forte de interesse, como assistir um vídeo longo do início ao fim, recebe mais peso que um indicador fraco. Nem a contagem de seguidores nem o desempenho anterior da conta são fatores diretos da recomendação.

Quem entra no app pode escolher categorias de interesse; quem não escolhe começa com um feed generalizado de vídeos populares. Cada nova interação ajuda o sistema a aprender os interesses do usuário.

Há também regras de variedade e elegibilidade: o feed geralmente não mostra dois vídeos seguidos com o mesmo som ou do mesmo criador e não recomenda conteúdo duplicado, já visto ou spam. Vídeos recém-enviados ou em revisão, e spam, também podem ser inelegíveis para recomendação.

Lidas em conjunto, as duas fontes cobrem camadas diferentes: o paper mostra como montar um treino online com tabela de embeddings sem colisão, e a newsroom mostra que sinais e regras de produto entram no ranqueamento. Este caso não liga uma à outra além disso.