Reconstrução de balanceamento vs Autoencoder variacional de perda de KL

Estou treinando um autoencoder variacional condicional em um conjunto de dados de faces. Quando eu defino minha Perda de KLL igual ao meu termo de Perda de reconstrução, meu autoencoder parece incapaz de produzir amostras variadas. Eu sempre recebo os mesmos tipos de rosto aparecendo:

Essas amostras são terríveis. No entanto, quando diminuo o peso da perda de KLL em 0,001, recebo amostras razoáveis:

O problema é que o espaço latente aprendido não é suave. Se eu tentar executar uma interpolação latente ou gerar uma amostra aleatória, fico com lixo. Quando o termo KLL tem um peso pequeno (0,001), observo o seguinte comportamento de perda: Observe que o VggLoss (o termo de reconstrução) diminui, enquanto o KLLoss continua a aumentar.

Também tentei aumentar a dimensionalidade do espaço latente, mas isso também não funcionou.

Observe aqui, quando os dois termos de perda têm peso igual, como o termo KLL domina, mas não permite que a perda de reconstrução diminua:

Isso resulta em terríveis reconstruções. Existem sugestões sobre como equilibrar esses dois termos de perda ou qualquer outra coisa possível a tentar, para que o meu autoencoder aprenda um espaço latente interpolativo e suave enquanto produz reconstruções razoáveis?

Um pouco tarde para a festa aqui e você provavelmente já passou disso, mas está bem documentado que você precisa "aquecer" o termo de perda de KL começando em zero e treinando um pouco apenas na perda de reconstrução antes de introduzir a perda ou os resultados de KL não seja bom. Não está claro no seu post se você fez isso, mas é um exemplo clássico de como esses treinamentos podem ser delicados - às vezes me pergunto como eles são apropriados na validação cruzada, já que são fortes tentativas e erros, além de um pouco de poeira de duende e arco-íris.

No entanto, quando diminuo o peso da perda de KLL em 0,001, recebo amostras razoáveis: (...) O problema é que o espaço latente aprendido não é suave.

Parece super ajuste. Lembre-se de que a perda de KL no espaço latente corresponde à regularização.

Existem sugestões sobre como equilibrar esses dois termos de perda ou qualquer outra coisa possível a tentar, para que meu autoencoder aprenda um espaço latente interpolável e suave, enquanto produz reconstruções razoáveis?

Recentemente me deparei com este artigo: $\beta$-VAE: Aprendendo conceitos visuais básicos com uma estrutura variacional restrita (na verdade, usa seu conjunto de dados em um exemplo).

Do papel ($\beta$ é o parâmetro que você alterou):

Introduzimos um hiperparâmetro ajustável $\beta$ que equilibra a capacidade do canal latente e as restrições de independência com a precisão da reconstrução (...) $\beta$-VAE é estável para treinamento, faz poucas suposições sobre os dados e depende do ajuste de um único hiperparâmetro $\beta$, que pode ser otimizado diretamente por meio de uma pesquisa de hiperparâmetro usando dados fracamente rotulados ou através de inspeção visual heurística para dados puramente não supervisionados.