Mastering memorization in flow-based generative models via the Moreau Envelope

Flow-based models have shown unprecedented generation capabilities across a wide range of modalities, but they suffer from the not fully understood phenomenon of memorization - the verbatim reproduction of data samples despite valid model convergence. This degrades their creativity and applicability in low data regimes as well as their reliability in secure settings. On the other hand, memorization could be the key to a new form of controlled generation, by memorizing features corresponding to specific states to observe, but this has not yet been studied for this purpose. We thus introduce a novel framework to achieve controlled memorization, with a dial allowing model optimization to decide to what extent it should be mitigated or enforced. This is made possible by turning a memorization metric into a Kantorovich potential via the Moreau Envelope, leveraging the proximal of the envelope to build an attractive or repulsive memorization field. We further provide an Optimal Transport interpretation of the proposed method, and derive a proximal gradient descent algorithm for the computation of the proximal during model inference. Moreover, we propose a number of illustrative use cases across various modalities; these range from preventing memorization in low data regime image generation and leverage the perturbation in text domain to steer flow-based language models token generation towards using specific words, showcasing the versatility of the proposed method.