# Neural-Kernel Conditional Mean Embeddings

arXiv: https://arxiv.org/abs/2403.10859 · HTML: https://mlsanity.com/paper/2403.10859/

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Placed in RL theory and sample complexity, one of 15 directions in Deep reinforcement learning methods.
- Builds on 19 of its 48 references that are in the map.
- First posted on arXiv 16 Mar 2024.

## Placed in

- Deep reinforcement learning methods / [RL theory and sample complexity](https://mlsanity.com/deeprl/rl-theory-and-sample-complexity/)

## Builds on (most cited first, 19 of 19)

- [Variational Inference with Normalizing Flows](https://mlsanity.com/paper/1505.05770/) (2015)
- [Deep Reinforcement Learning with Double Q-learning](https://mlsanity.com/paper/1509.06461/) (2015)
- [Spectral Normalization for Generative Adversarial Networks](https://mlsanity.com/paper/1802.05957/) (2018)
- [A Distributional Perspective on Reinforcement Learning](https://mlsanity.com/paper/1707.06887/) (2017)
- [Neural Spline Flows](https://mlsanity.com/paper/1906.04032/) (2019)
- [Distributional Reinforcement Learning with Quantile Regression](https://mlsanity.com/paper/1710.10044/) (2017)
- [Gaussian Processes for Big Data](https://mlsanity.com/paper/1309.6835/) (2013)
- [Kernel Mean Embedding of Distributions: A Review and Beyond](https://mlsanity.com/paper/1605.09522/) (2016)
- [Benchmarking Simulation-Based Inference](https://mlsanity.com/paper/2101.04653/) (2021)
- [Kernel Instrumental Variable Regression](https://mlsanity.com/paper/1906.00232/) (2019)
- [Overcoming Limitations of Mixture Density Networks: A Sampling and Fitting Framework for Multimodal Future Prediction](https://mlsanity.com/paper/1906.03631/) (2019)
- [Revisiting Rainbow: Promoting more Insightful and Inclusive Deep Reinforcement Learning Research](https://mlsanity.com/paper/2011.14826/) (2020)
- [Learning Deep Features in Instrumental Variable Regression](https://mlsanity.com/paper/2010.07154/) (2020)
- [normflows: A PyTorch Package for Normalizing Flows](https://mlsanity.com/paper/2302.12014/) (2023)
- [Conditional Distributional Treatment Effect with Kernel Conditional Mean Embeddings and U-Statistic Regression](https://mlsanity.com/paper/2102.08208/) (2021)
- [MMD-FUSE: Learning and Combining Kernels for Two-Sample Testing Without Data Splitting](https://mlsanity.com/paper/2306.08777/) (2023)
- [Counterfactual Mean Embeddings](https://mlsanity.com/paper/1805.08845/) (2018)
- [BayesIMP: Uncertainty Quantification for Causal Data Fusion](https://mlsanity.com/paper/2106.03477/) (2021)
- [Bayesian Learning of Conditional Kernel Mean Embeddings for Automatic Likelihood-Free Inference](https://mlsanity.com/paper/1903.00863/) (2019)

## Built on by (most cited first, 7 of 7)

- [Foundations of Multivariate Distributional Reinforcement Learning](https://mlsanity.com/paper/2409.00328/) (2024)
- [Optimal Convergence Rates for Neural Operators](https://mlsanity.com/paper/2412.17518/) (2024)
- [An Overview of Causal Inference using Kernel Embeddings](https://mlsanity.com/paper/2410.22754/) (2024)
- [Tensor-Var: Efficient Four-Dimensional Variational Data Assimilation](https://mlsanity.com/paper/2501.13312/) (2025)
- [Conditional Distributional Treatment Effects: Doubly Robust Estimation and Testing](https://mlsanity.com/paper/2603.16829/) (2026)
- [The Generalised Kernel Covariance Measure](https://mlsanity.com/paper/2604.03721/) (2026)
- [Multi-Task Learning of Conditional Mean Operators: applications to dynamical systems and uncertainty quantification](https://mlsanity.com/paper/2609.35429/) (2026)

## Abstract (arXiv)

Kernel conditional mean embeddings (CMEs) offer a powerful framework for representing conditional distribution, but they often face scalability and expressiveness challenges. In this work, we propose a new method that effectively combines the strengths of deep learning with CMEs in order to address these challenges. Specifically, our approach leverages the end-to-end neural network (NN) optimization framework using a kernel-based objective. This design circumvents the computationally expensive Gram matrix inversion required by current CME methods. To further enhance performance, we provide efficient strategies to optimize the remaining kernel hyperparameters. In conditional density estimation tasks, our NN-CME hybrid achieves competitive performance and often surpasses existing deep learning-based methods. Lastly, we showcase its remarkable versatility by seamlessly integrating it into reinforcement learning (RL) contexts. Building on Q-learning, our approach naturally leads to a new variant of distributional RL methods, which demonstrates consistent effectiveness across different environments.