FedMentor: Domain-Aware Differential Privacy for Heterogeneous Federated LLMs in Mental Health

Resumo Técnico: FedMentor

Declaração do Problema

A implantação de Grandes Modelos de Linguagem (LLMs) para aplicações de saúde mental enfrenta um trio crítico de desafios: requisitos rigorosos de privacidade (ex: HIPAA, GDPR), a heterogeneidade dos dados dos usuários em diferentes domínios clínicos e as restrições computacionais de ambientes de saúde com recursos limitados. O treinamento centralizado tradicional é frequentemente inviável devido à sensibilidade dos dados de saúde mental, que podem incluir ideação de autolesão explícita ou sinais clínicos que exigem confidencialidade robusta. Além disso, as abordagens de Aprendizado Federado (FL) existentes frequentemente lutam contra desvios de dados não-IID (não independentes e não identicamente distribuídos) entre domínios e falham em fornecer garantias de privacidade granulares adaptadas a sensibilidades de dados específicas. Há uma necessidade de um framework que permita o ajuste fino colaborativo de LLMs enquanto preserva a privacidade do usuário, mantém a utilidade do modelo e garante a eficiência de comunicação em dispositivos de borda (edge devices).

Metodologia: FedMentor

Os autores propõem o FedMentor, um framework de ajuste fino federado projetado especificamente para domínios heterogêneos de saúde mental. O framework integra três componentes principais:

1. Adaptação de Baixo Posto Federada (LoRA): Em vez de transmitir os pesos totais do modelo, o FedMentor congela o backbone do LLM global e treina e comunica apenas adaptadores LoRA leves. Isso reduz significativamente a sobrecarga de comunicação e os requisitos de memória, permitindo o treinamento de modelos com até 1,7 bilhão de parâmetros em clientes de GPU única.
2. Privacidade Diferencial Consciente do Domínio (DP): Reconhecendo que diferentes conjuntos de dados de saúde mental possuem níveis variados de sensibilidade, o FedMentor atribui orçamentos de privacidade ($\epsilon, \delta$) distintos para cada domínio de cliente.
   • Alocação de Orçamento: Domínios de alta sensibilidade (ex: Fatores de Risco Interpessoal) recebem orçamentos mais rigorosos ($\epsilon$ menor), enquanto domínios de menor sensibilidade recebem orçamentos mais amplos.
   • Calibração de Ruído: O framework emprega um mecanismo de injeção de ruído consciente da camada e do adaptador, escalonando o ruído Gaussiano com base na camada da rede (inicial vs. tardia) e na matriz LoRA específica (A vs. B), reconhecendo que as camadas iniciais e as matrizes LoRA-A são mais sensíveis à perturbação.
3. Loop de Feedback de Utilidade-Privacidade Adaptativo: Para evitar que o ruído excessivo de privacidade degrade o desempenho do modelo, o servidor monitora métricas de utilidade de proxy (ex: BERTScore, Taxa de Segurança). Se a utilidade agregada cair abaixo de um limiar predefinido ($\tau$), o servidor reduz adaptativamente a escala do ruído para rodadas subsequentes, negociando dinamicamente a rigidez da privacidade pela utilidade clínica.

O pipeline de treinamento envolve clientes treinando adaptadores LoRA locais em backbones congelados, adicionando ruído Gaussiano específico do domínio às atualizações e transmitindo esses adaptadores ruidosos para um servidor central. O servidor agrega as atualizações usando um algoritmo FedAvg ponderado pelo conjunto de dados.

Principais Contribuições

O artigo descreve três contribuições primárias:

• Adaptação Privada Consciente do Domínio: A formulação de uma abordagem de LoRA federada com orçamentos de privacidade por domínio que protege textos sensíveis de saúde mental enquanto preserva a utilidade da tarefa. Esta abordagem eleva as métricas de segurança (ex: aumento da Taxa de Segurança de Toxicidade) enquanto mantém as métricas de utilidade (BERTScore F1, ROUGE-L) próximas das bases não privadas.
• Robustez à Heterogeneidade Não-IID: A demonstração de utilidade estável e equidade entre diversos conjuntos de dados (Dreaddit, IRF, MultiWD) sem colapso de domínio. O mecanismo de ruído adaptativo garante que o desempenho seja mantido mesmo sob orçamentos de privacidade rigorosos para os domínios mais sensíveis.
• Eficiência Prática via Adaptadores: Ao trocar apenas os pesos do adaptador LoRA, o framework alcança uma compressão significativa (reduzindo os custos de comunicação por um fator de $\approx 10^{-3}$), tornando viável o treinamento de modelos de até 1,7B de parâmetros em clientes com limitações de GPU única.

Resultados Experimentais

O framework foi avaliado em três conjuntos de dados de saúde mental (Dreaddit, IRF, MultiWD) e cinco backbones de LLM variando de 350M a 1,7B de parâmetros.

• Segurança e Privacidade: O FedMentor melhorou a segurança em relação ao Aprendizado Federado padrão sem privacidade. No conjunto de dados MultiWD, a Taxa de Segurança de Toxicidade (TSR) aumentou 2 pontos percentuais, e a toxicidade média diminuiu de 2,92 para 1,98.
• Preservação de Utilidade: Apesar da adição de privacidade diferencial, as métricas de utilidade permaneceram dentro de 0,5% da base FL não privada e próximas do limite superior centralizado. Por exemplo, as pontuações BERTScore F1 e ROUGE-L mostraram degradação negligenciável.
• Eficiência: O framework treinou com sucesso modelos de até 1,7B de parâmetros em clientes de GPU única. Os custos de comunicação foram reduzidos para aproximadamente 173 MB por rodada para o maior modelo, com o uso de memória de pico permanecendo dentro dos limites de GPUs de 80 GB.
• Estudos de Ablação: Experimentos confirmaram que orçamentos específicos de domínio e ruído adaptativo são críticos. Orçamentos de privacidade uniformes melhoraram a equidade nos resultados de segurança, enquanto cronogramas de ruído estáticos levaram à redução da segurança e ao aumento da disparidade entre os clientes.

Significância e Alegações

O artigo afirma ser o primeiro estudo a combinar Privacidade Diferencial consciente do domínio com ajuste fino de LoRA federado especificamente para LLMs em saúde mental. Os autores posicionam o FedMentor como uma solução prática para o "trade-off privacidade-utilidade" que assola as implementações atuais de IA em saúde mental. Ao permitir o aprendizado colaborativo sem centralizar dados sensíveis, o framework visa aumentar a confiança e a aplicabilidade clínica em domínios de alto risco.

Os autores mantêm uma postura modesta em relação às suas alegações, reconhecendo limitações como o uso de conjuntos de dados de redes sociais apenas em inglês (tratando cada um como um único cliente), a dependência de métricas de segurança automatizadas em vez de revisão clínica e a ausência de auditorias de subgrupos demográficos. Eles observam que o estudo reflete um ambiente de hardware específico (NVIDIA A100) e que trabalhos futuros focarão em escalar federações, incorporar abordagens com participação do clínico (clinician-in-the-loop) e realizar avaliações multilíngues.