CRUMB: Efficient Prior Fitted Network Inference via Distributionally Matched Context Batching

Resumo Técnico: CRUMB – Inferência de Redes Ajustadas a Priores via Lote de Contexto Correspondido Distribuicionalmente

Definição do Problema
Redes ajustadas a priores (PFNs - Prior-Fitted Networks), como o TabPFN e o TabICL, representam uma classe de modelos de fundação tabulares capazes de aprendizado em contexto (in-context learning). Elas aceitam um conjunto de treinamento totalmente rotulado como contexto e geram previsões para consultas de teste em uma única passagem direta (forward pass), frequentemente superando métodos tradicionais de gradient boosting em conjuntos de dados pequenos a médios. No entanto, um gargalo crítico limita sua escalabilidade: o mecanismo de autoatenção em transformadores baseados em PFNs escala quadraticamente ($O(N^2)$) com o número de amostras de treinamento ($N$). Isso torna a inferência computacionalmente proibitiva para grandes conjuntos de dados ($N > 50.000$).

Soluções existentes enfrentam um compromisso entre precisão e eficiência. A subamostragem uniforme melhora a velocidade ao descartar aleatoriamente dados de treinamento, mas sacrifica informação. Métodos dependentes de consulta, como os $k$-vizinhos mais próximos ($k$NN), melhoram a precisão ao adaptar o contexto para cada ponto de teste, mas impedem o processamento em lote (batching), exigindo $T$ passagens diretas separadas para $T$ pontos de teste. Outras abordagens, como o Mixture of In-Context Prompters (MICP), tentam agrupar consultas por meio de agrupamento (clustering) de dados de treinamento, mas podem falhar em alinhar o contexto de treinamento com a distribuição específica dos dados de teste, particularmente sob deriva de covariáveis (covariate drift).

Metodologia
Os autores propõem o CRUMB (Clustered Retrieval Using Minimised-MMD Batching), um invólucro (wrapper) de inferência agnóstico à arquitetura e de três estágios que não requer o retreinamento do modelo. O método resolve a tensão entre qualidade do contexto e eficiência de lote através dos seguintes estágios:

1. Agrupamento do Lado do Teste: As consultas de teste ($D_{test}$) são particionadas em $K$ clusters ($C_1, \dots, C_K$) usando $k$-means em características de entrada padronizadas. Isso permite que o sistema agrupe pontos de teste similares que podem compartilhar um contexto de treinamento comum.
2. Seleção de Contexto Correspondido Distribuicionalmente: Para cada cluster de teste $C_k$, um pequeno subconjunto de treinamento $S_k$ (tamanho $n \ll N$) é selecionado do pool total de treinamento ($D_{train}$). Diferente de métodos de proximidade geométrica (ex: vizinho mais próximo do centroide), o CRUMB seleciona $S_k$ minimizando de forma gananciosa a Discrepância de Média Máxima (MMD - Maximum Mean Discrepancy) entre a distribuição empírica do cluster de teste e o subconjunto de treinamento candidato.
   • Isso é alcançado via kernel herding, um algoritmo ganancioso que seleciona iterativamente pontos de treinamento para minimizar o objetivo MMD. A função objetivo equilibra dois termos: recompensar a proximidade com o cluster de teste (relevância) e penalizar a redundância entre os pontos de treinamento selecionados (diversidade).
   • Para acelerar esse processo, os autores empregam Características de Fourier Aleatórias (RFF - Random Fourier Features) para aproximar o kernel RBF Gaussiano e utilizam seleção gananciosa em lote.
3. Inferência em Lote: O modelo PFN executa $K$ passagens diretas independentes. Em cada passagem, o modelo recebe o subconjunto de treinamento selecionado $S_k$ como contexto e o cluster inteiro $C_k$ como consultas. Isso reduz o número total de passagens diretas de $T$ (no caso do $k$NN por consulta) para $K$, permitindo a utilização eficiente da GPU enquanto mantém o alinhamento distribuicional.

Principais Contribuições

• Recuperação de Contexto Baseada em MMD: O artigo introduz a minimização gananciosa de MMD (kernel herding) como um mecanismo para selecionar subconjuntos de treinamento. Os autores demonstram que essa abordagem de correspondência distribuicional supera as estratégias de vizinho mais próximo de centroide e de recuperação uniforme de Voronoi dentro do mesmo framework de agrupamento.
• Seleção de Contexto em Lote via Agrupamento do Lado do Teste: Ao agrupar consultas de teste e atribuir um único contexto de treinamento, correspondido distribuicionalmente, a cada cluster, o CRUMB permite um processamento em lote eficiente sem sacrificar a relevância da consulta.
• Robustez à Deriva de Covariáveis: O método lida naturalmente com a deriva de covariáveis. Como o contexto de treinamento é selecionado para corresponder à distribuição atual do lote de teste (em vez de uma distribuição de treinamento fixa), o CRUMB se adapta a mudanças nos dados de teste. Em contraste, métodos como o MICP, que dependem de clusters de treinamento fixos, degradam significativamente quando a distribuição de teste se desloca da distribuição de treinamento.

Resultos Experimentais
Os autores avaliaram o CRUMB no benchmark TabArena (51 conjuntos de dados tabulares diversos) em três arquiteturas de PFN: TabPFNv2, TabICLv1 e TabICLv2.

• Desempenho vs. Baselines: No benchmark TabArena, o CRUMB supera significativamente a subamostragem uniforme e a técnica MICP com orçamentos de contexto idênticos. Embora o CRUMB seja estatisticamente próximo em desempenho ao $k$NN por consulta, ele o alcança com um número fixo de $K$ passagens diretas (ex: $K=20$) em vez de $T$ passões, resultando em uma redução substancial no tempo de inferência (ex: 35s vs. 361s nos experimentos relatados).
• Grandes Conjuntos de Dados: Experimentos em conjuntos de dados com até 120.000 amostras (ex: GiveMeSomeCredit, Higgs) mostram que o CRUMB com parada antecipada (early stopping — terminando a seleção de contexto quando a melhoria do MMD estagna) vence com mais frequência o MICP, particularmente para TabPFNv2 e TabICLv1.
• Deriva de Covariáveis: Em experimentos de deriva controlada onde a distribuição de teste se desloca ao longo do primeiro componente principal, a vantagem do CRUMB sobre o MICP cresce monotonicamente com a intensidade da deriva. Com zero deriva, o CRUMB lidera com +4,9%; na deriva máxima (completamente fora da amostra), a lacuna aumenta para +17,1%. Isso é atribuído à capacidade do CRUMB de realinhar o contexto de treinamento à distribuição de teste deslocada, enquanto os clusters de treinamento fixos do MICP tornam-se desalinhados.

Significância e Alegações
O artigo afirma que o CRUMB torna a inferência de redes ajustadas a priores prática para conjuntos de dados tabulares de grande escala sem exigir mudanças arquitetônicas ou ajuste fino (fine-tuning) do modelo. Ao combinar o agrupamento do lado do teste com a recuperação de contexto correspondido distribuicionalmente, o CRUMB atinge um "ponto ideal" entre a precisão da recuperação por consulta e a eficiência do processamento em lote uniforme.

Os autores enfatizam que a inovação central reside no design adaptável ao teste: agrupar pontos de teste e selecionar contextos baseados no alinhamento distribuicional (MMD) em vez de apenas proximidade geométrica. Esse design fornece um mecanismo natural de robustez contra a deriva de covariáveis, um regime onde os métodos de seleção de contexto existentes frequentemente falham. O artigo conclui que, embora o CRUMB herde o teto de desempenho do pré-treinamento do PFN subjacente, ele remove efetivamente a barreira computacional que impede a aplicação de PFNs a grandes conjuntos de dados.

Limitações
Os autores reconhecem que o CRUMB requer o conjunto completo de consultas de teste antecipadamente para realizar o agrupamento, tornando-o inadequado para cenários de fluxo contínuo (streaming) reais sem modificações (ex: cache de centroides, o que efetivamente reduz ao roteamento estilo MICP). Além disso, a dependência de $k$-means euclidiano e kernels Gaussianos isotrópicos pode limitar o desempenho em espaços de características de alta dimensão ou tipos mistos. Finalmente, como um invólucro de inferência puro, o CRUMB não pode compensar um casamento ruim entre os priors de pré-treinamento do PFN e o conjunto de dados específico.