Flow Density Control: Generative Optimization Beyond Entropy-Regularized Fine-Tuning
Optimizing functionals of distributions beyond entropy and expectation.
Marin Vlastelica
Hi, I am a member of technical staff at Project Prometheus.
Academically, I was a postdoctoral researcher at the ETH AI Center, under the supervision of Andreas Krause, Marco Hutter and Bernhard Schoelkopf. I completed my PhD degree summa cum laude, examined by Georg Martius, Peter Dayan, Mathias Niepert and Matthias Hein, at the Max Planck Institute for Intelligent Systems in Tuebingen, Germany as part of the International Max Planck Research School for Intelligent Systems. My thesis advisory committee consisted of Georg Martius, Peter Dayan, and Michael Muehlebach, with my primary supervisor being Georg Martius. In addition, I hold a master's degree in computer science with a machine learning focus from the Karlsruhe Institute of Technology, Germany, and a bachelor's degree in computer science from the University of Zagreb, Croatia.
In 2017 I was fortunate to receive a research fellowship with the Blue Brain Project in Geneva, Switzerland, where I came into contact with neuroscience. In 2021 I spent some time at Amazon as an applied scientist intern working on variational inference, reinforcement learning and NLP in collaboration with Patrick Ernst and Gyuri Szarvas. I spent part of 2022 at DeepMind, hosted by Kimberly Stachenfeld in Peter Battaglia's group, working on diffusion generative modelling for building better optimizers in close collaboration with Arnaud Doucet. In 2024 I was visiting Meta GenAI Zurich, where I worked on building more efficient large-scale diffusion generative models. Most recently, I was a research scientist at Meta Super Intelligence Labs working on post-training.
My research interests revolve around principled methods for discovery and self-improvement — creating intelligence that can continually adapt and become more efficient in learning and extracting new knowledge. To this end, I have gained rich experience in several sub-areas of machine learning such as reinforcement learning and large generative models, and have led several efforts in academic and industrial settings. Some of my more recent efforts are centered on post-training for agentic LLMs and distributional optimization in diffusion and flow-matching models.
In my free time I enjoy playing guitar, reading opinionated books on various topics, and all sorts of sports activities.
Feel free to get in touch:
Research
Provable maximum entropy manifold exploration via diffusion models
Controllable distributional manifold exploration.
We provide a principle algorithm for maximum entropy constrained exploration on the distributional manifold.
Diverse Offline Imitation via Fenchel Duality
Leveraging duality for skill extraction
We show that Fenchel duality can be utilized to maximize a diversity objective subject to f-divergence constraints.
Diffusion Generative Inverse Design
Building better optimizers by guided diffusion
We propose an improved sampling procedure for a diffusion generative model for approximate sampling from a target distribution that is defined by an energy (cost) function of designs in a complex simulation environment.
Learning Agile Skills via Adversarial Imitation of Rough Partial Demonstrations
Wasserstein adversarial learning of 'bad' demonstrations
In this work we proposed a method for imitating rough demonstrations by a robot quadruped based on Wasserstein adversarial learning of an imitation reward. We trained a policy via domain randomization and successfully transferred it to the real system, where to robot was able to reproduce agile motions.
Gradient Backpropagation Through Combinatorial Algorithms: Identity with Projection Works
No need for a solver call on the backward pass (sometimes)
As a continuation of the blackbox-backprop line of work, we introduce a simple modification to the algorith, namely treating the solver as an identity mapping in the computation graph on the backward pass.
This, coupled with projections that avoid degenerate cases, works comparably well as blackbox-backprop.
Versatile Skill Control via Self-supervised Adversarial Imitation of Unlabeled Mixed Motions
Extracting meaningful behaviors by adversarial learning.
Adversarial learning to facilitate unsupervised skill discovery resulting in a controllable skill set by a latent variable with skills that are useful for solving the downstream task.
Risk-Averse Zero-Order Trajectory Optimization
How can we account for uncertainty in zero-order trajectory optimizers?
In this work we address the problem of efficient uncertainty estimation in zero-order trajectory optimization via the Cross Entropy Method (CEM). Moreover, we show that it is essential that we can distinguish between epistemic and aleatoric uncertainty in order to avoid so-called “risky” behaviour.
Taming Continuous Posteriors for Latent Variational Dialogue Policies
Making continuous latents for task-oriented dialogue great again
Latent action reinforcement learning for task-oriented dialogue has seen success at benchmarks such as MultiWOZ. Categorical latents have been argued to be the best choice. We show that with continuous latents and reformulation of the ELBO objective and the reinforcmenet learning stage, we can achieve state-of-the-art performance on MultiWOZ.
Neuro-Algorithmic Policies Enable Fast Combinatorial Generalization
Introducing neuro-algorithmic policy architecture.
As a continuation of the blackbox-differentiation line of work, we propose to use time-dependent shortest-path solvers in order to enhance generalization capabilities of neural network policies. With imposing a prior on the underlying goal-conditioned MDP structure, we are able to extract well-performing policies through imitation learning that utilize blackbox solvers for receding horizon planning at execution time. Again, this comes with absolutely no sacrifices to the optimality of the solver used.
Optimizing Rank-based Metrics via Blackbox Differentiation
Ranking as a blackbox combinatorial problem.
As another continuation of the blackbox differentiation line of work, we show that we can cast the ranking problem as a blackbox solver that satisfies the conditions for efficient gradient calculation, therefore enabling us to optimize rank-based metrics by simply using efficient implementations of sorting algorithms instead of learning a differentiable sort operation. We apply this insight to optimizing mean average precision and recall in object detection and retrieval tasks, where we achieve comparable results to state-of-the-art at the time.
Differentiation of Blackbox Combinatorial Solvers
Can we embed combinatorial solvers in neural architectures?
Problems that are inherently combinatorial still remain a hinderance for classical deep learning methods. Traditional methods that try to do gradient propagation through combinatorial solvers rely on sample-based estimates or solver relaxations. We show that for a specific class of solver, we are able to efficiently compute gradients of an implicit piecewise-linear interpolation of the objective. This allows us to achieve unprecedented generalization performance on representation learning tasks with combinatorial flavor.
Control What You Can: Intrinsically Motivated Task-Planning Agent
Extracting planning graphs by intrinsic motivation.
In this work we propose a hierarchical reinforcement learning algorithm that is able to construct planning graphs while managing to use computational resources efficiently, guided by intrinsic motivation in form of prediction error and measure of task improvement.
Reinforcement Learning with Liquid State Machines
Can we use spiking neural networks for continuous control?
In this work we propose a liquid state machine approach to reinforcement learning of continuous motor control. The liquid state machine approach with smart spectral radius initialization has shown to extract useful features for motor control which can be used with classical RL algorithms.
KIT at MediaEval 2015-Evaluating Visual Cues for Affective Impact of Movies Task
Evolutionary methods for minimal addition chain exponentiation
In this work we looked at what kind of features we can use for prediction of video valence, violence and sentiment.
Side Projects
FOX (FlOws in JAX)
Normalizing Flows in JAX
Developing a Python package with standard normalizing flow implementations that is based on JAX and Flax.
Decentralized Energy Exchange
Based on Ethereum blockchain
As part of hackathon, we have noticed that the future of energy distribution networks are with small producers (via solar cells for example). The missing component is an exchange that makes it possible for producers to exchange energy for value, which we tried to make possible. With this, we have won 1st place on a Microsoft-organized hackathon Germany-wide.