Entity linking

01 · CONTEXTO

O problema: ligar menções a entidades

Entity linking recebe um texto e a lista de menções de entidades que ele contém, e devolve uma lista de pares menção-entidade, em que cada entidade é uma entrada de uma base de conhecimento, como a Wikipedia.

Na lista de exercícios de projeto de sistemas de ML de Chip Huyen aparecem perguntas desse tipo: treinar um modelo para prever se uma palavra numa frase se refere ao animal ou ao carro, e mapear apelidos como Pete, Andy, Nick e Rob para nomes reais.

O que decide a entidade é o contexto da frase, não a grafia da menção: a mesma palavra aponta para entradas diferentes da base conforme o texto ao redor.

A própria página avisa que muitas perguntas são ambíguas de propósito e que cabe a quem é entrevistado pedir esclarecimento e delimitar o escopo do problema. O curso Grokking the Machine Learning Interview, do Educative, lista um projeto Entity Linking System construído sobre linking de entidades nomeadas (NEL), com reconhecimento, desambiguação, métricas, arquitetura e modelagem; só essa descrição pública foi usada aqui.

02 · FORMULAÇÃO

O que o paper de referência assume

Como desenho de referência, o caso usa o paper Scalable Zero-shot Entity Linking with Dense Entity Retrieval, do Facebook AI Research. Ele estuda o linking com menções já fornecidas (gold mentions) e assume que título e descrição das entidades estão disponíveis.

O desafio central é a escala: para cada menção há milhões de entidades possíveis a considerar. Métodos anteriores recorriam a tabelas de menções e à popularidade de links da Wikipedia; no cenário zero-shot do paper, não há conhecimento externo e cada entidade é definida apenas por uma descrição textual curta.

03 · ARQUITETURA

Recuperação densa seguida de reranking

O paper divide a tarefa em duas etapas: geração de candidatos e ranqueamento. Trabalhos anteriores geravam candidatos com informação de frequência, tabelas de alias e métodos baseados em TF-IDF.

A proposta é um algoritmo zero-shot em dois estágios. No primeiro, a recuperação ocorre num espaço denso definido por um bi-encoder que embute de forma independente o contexto da menção e as descrições das entidades. No segundo, cada candidato é reranqueado por um cross-encoder que concatena o texto da menção e o da entidade.

1. Fora do caminho da requisição, um job em lote lê título e descrição de cada entidade na base de conhecimento, calcula seus embeddings e os publica no índice vetorial. As representações das entidades podem ser pré-computadas e guardadas em cache.

2. O cliente envia ao serviço de linking o texto e a menção a resolver.

3. O serviço pede ao codificador de menção o vetor da menção com seu contexto; o bi-encoder embute menção e entidades de forma independente.

4. O índice devolve as entidades de maior produto escalar com o vetor da menção.

5. O reranker lê as descrições dos candidatos na base de conhecimento e pontua cada par menção-entidade; o cross-encoder tem atenção cruzada profunda entre o contexto e as descrições.

6. O serviço devolve ao cliente a entidade escolhida para a menção.

04 · MODELOS E TREINO

Bi-encoder, cross-encoder e como treinar

No bi-encoder, dois transformers BERT independentes codificam o contexto com a menção e a entidade em vetores densos, e cada candidato é pontuado pelo produto escalar desses vetores. A entidade entra com título e descrição; para entidades da Wikipedia, a descrição usa as primeiras frases da página.

No cross-encoder, contexto/menção e entidade são codificados juntos num só transformer, com uma camada linear adicional. Por causa da maior pegada de memória e computação, ele é usado no reranking de um conjunto pequeno de candidatos do bi-encoder, até 100, e o paper o considera inadequado para recuperação ou tarefas que exigem inferência rápida.

O bi-encoder é treinado para maximizar o score da entidade correta frente às entidades amostradas aleatoriamente do mesmo lote. Além dos negativos do lote, o paper usa negativos difíceis, obtidos com as 10 entidades mais bem previstas para cada exemplo de treino e somados aos negativos aleatórios.

No experimento zero-shot, o cross-encoder foi treinado sobre os 64 primeiros candidatos do bi-encoder para cada amostra, com perda softmax para maximizar o score da entidade correta.

1. Os pares menção-entidade anotados alimentam o treino do bi-encoder, que maximiza o score da entidade correta frente às entidades do mesmo lote.

2. A mineração escolhe, por exemplo de treino, as entidades mais bem previstas como negativos difíceis e as devolve ao treino, junto dos negativos aleatórios do lote.

3. O bi-encoder treinado fornece os candidatos para o treino do cross-encoder: no experimento zero-shot, os 64 primeiros de cada amostra.

4. O cross-encoder é treinado com perda softmax para maximizar o score da entidade correta entre os candidatos.

5. Na destilação, o cross-encoder faz o papel de professor de um bi-encoder: parte do ganho de acurácia do cross-encoder, mais caro, é transferida ao bi-encoder.

6. Os modelos resultantes são gravados no repositório de modelos.

05 · AVALIAÇÃO E CUSTO

Recall, latência e o preço de cada estágio

Na avaliação zero-shot, o bi-encoder atinge recall bem maior que o do BM25, e o paper passa os 64 primeiros candidatos recuperados ao ranqueador.

Em velocidade, o paper afirma que o linking com bi-encoder é muito rápido com busca aproximada de vizinhos, por exemplo sobre 5.9 milhões de candidatos em 2 milissegundos. Calcular os embeddings do conjunto de 5.9M entidades é um passo intensivo em recursos, mas paralelizável, e levou cerca de 2.8 horas.

O paper também mostra que boa parte do ganho de acurácia do cross-encoder, mais caro, pode ser transferida ao bi-encoder por destilação de conhecimento.

06 · LIMITES E ENTREVISTA

O que o desenho não cobre

O paper foca em resolução local, em que cada menção é modelada de forma independente. A predição fora da base (nil) e o linking de ponta a ponta ficam para trabalho futuro: o paper diz que suas técnicas deveriam se aplicar a esse cenário, mas deixa a exploração para depois.

Para a entrevista, depois de fechar o escopo (menções dadas ou detectadas, tamanho da base, o que fazer sem entidade correta), apresente recuperação barata seguida de reranking caro, justifique o pré-cálculo dos embeddings das entidades, discuta negativos difíceis no treino e termine pelos limites: coerência entre menções do mesmo documento, entidades novas e menções sem entidade na base.