Research map / Reinforcement learning / Visual and representation RL
Improving Unsupervised Hierarchical Representation with Reinforcement Learning
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Trains hierarchical VAEs with policy-gradient RL, using a skip-generative-path reward and a discounted Q-value so that each latent level is regularized by its mutual information with x, mitigating posterior collapse and improving unsupervised representations.
Problem
Hierarchical VAEs suffer from posterior collapse in higher-level latent variables, so they learn poor hierarchical representations for downstream tasks.
Approach
Formulates hierarchical VAE inference as a sequential decision process; a reward from a skip-generative path (marginalizing lower latents) measures information between x and each level, and a discounted Q-value is maximized directly by reparameterized policy gradient, avoiding REINFORCE-style estimators.
Key results
- HVAE+RL reaches 78.23% accuracy on FashionMNIST versus 10.73% for HVAE, and 46.46% on CIFAR10 versus 10.40%.
- LVAE+RL reaches 25.41% accuracy on CIFAR100 versus 10.38% for LVAE, and outperforms IVAE across all five datasets.
- Using RL-trained models, top-level-only representations (zL) are competitive with full-z representations (e.g., 78.48% vs 80.25% for LVAE on FashionMNIST) while the paper states that ELBO improvement is not required for better representations.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| apply RL policy-gradient training to HVAE | HVAE · FashionMNIST, CIFAR10, CIFAR100, TinyImageNet, STL-10 | accuracy: 10.73±0.08 -> 78.23±0.13; 10.40±0.05 -> 46.46±0.15; 1.21±0.10 -> 25.21±0.21; 0.63±0.06 -> 12.76±0.39; 10.24±0.07 -> 43.7 | Helped |
| apply RL policy-gradient training to LVAE | LVAE · FashionMNIST, CIFAR10, CIFAR100, TinyImageNet, STL-10 | accuracy: 19.70±0.21 -> 78.48±0.15; 15.32±0.10 -> 46.35±0.17; 10.38±0.89 -> 25.41±0.20; 4.77±0.47 -> 13.41±0.44; 11.89±0.10 -> 41. | Helped |
| use only the current-step reward instead of accumulated Q-value rewards (γ=0) (ablation) | HVAE+RL · FashionMNIST and CIFAR10 | accuracy: 75.63±0.11 vs 78.23±0.13; 41.91±0.15 vs 46.46±0.15 | Hurt |
| use only the current-step reward instead of accumulated Q-value rewards (γ=0) (ablation) | LVAE+RL · FashionMNIST and CIFAR10 | accuracy: 75.52±0.10 vs 78.48±0.15; 42.16±0.13 vs 46.35±0.17 | Hurt |
| vary the discount factor among γ=0.8, 0.85, 0.9, 0.95, and 0.98 (ablation) | HVAE+RL and LVAE+RL · FashionMNIST and CIFAR10 | accuracy: 78.03±0.16; 78.11±0.24; 78.23±0.13; 77.95±0.19; 78.01±0.10; 46.45±0.30; 45.89±0.24; 46.46±0.15; 46.21±0.21; 46.31±0.20 | No clear effect |
| increase the number of hierarchical levels (ablation) | HVAE, LVAE, IVAE, HVAE+RL, and LVAE+RL · FashionMNIST and CIFAR10 | accuracy | Mixed |
| use the full hierarchy z instead of only the topmost latent zL (ablation) | HVAE+RL and LVAE+RL · FashionMNIST and CIFAR10 | accuracy: 78.23±0.13 -> 84.56±0.25; 78.48±0.15 -> 84.77±0.30; 46.46±0.15 -> 52.87±0.54; 46.35±0.17 -> 52.89±0.48 | Helped |
Surprising findings
- Better ELBO/bpd does not correspond to better representation learning; BIVA and LVAE have strong bpd on CIFAR10 but weaker downstream accuracy than IVAE and the RL methods.
- Setting the discount factor to 0 (only immediate reward) lowers accuracy substantially compared to Q-value with gamma 0.9 (e.g., 75.63% vs 78.23% on FashionMNIST for HVAE).
Limitations
- Hierarchical VAE-based representation learning may be outdated; the approach is only demonstrated on VAEs and not on diffusion or other hierarchical models.
- Evaluation is limited to linear SVM probing on small image datasets with 5 seeds.
- Mutual information estimates via MINE show mixed results, e.g., LVAE+RL has lower I(x,zL) than LVAE-ELBO variants on some datasets like CIFAR10 relative to IVAE.
Experimental setup
Datasets: FashionMNIST CIFAR10 CIFAR100 TinyImageNet STL-10
Benchmarks: FashionMNIST CIFAR10 CIFAR100 TinyImageNet STL-10
Metrics: linear SVM classification accuracy MINE mutual information estimate I(x,zL) bits per dimension KL divergence at top level
Baselines: AAE VAE BIVA IVAE HVAE LVAE
Builds on: hierarchical VAE informative hierarchical VAE (IVAE) ladder VAE policy gradient
Closest papers by meaning
- Probabilistic Subgoal Representations for Hierarchical Reinforcement Learning (ICML 2024 · similarity 0.88)
- Structural Information-based Hierarchical Diffusion for Offline Reinforcement Learning (NeurIPS 2025 · similarity 0.87)
- Towards Principled Representation Learning from Videos for Reinforcement Learning (ICLR 2024 · similarity 0.86)
- Hierarchical Reinforcement Learning with Targeted Causal Interventions (ICML 2025 · similarity 0.86)
- Hierarchical Reinforcement Learning with Uncertainty-Guided Diffusional Subgoals (ICML 2025 · similarity 0.86)
- Online Laplacian-Based Representation Learning in Reinforcement Learning (ICML 2025 · similarity 0.86)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.