Research map / Reinforcement learning / Visual and representation RL
Learning to Scale Logits for Temperature-Conditional GFlowNets
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.
Proposes Logit-GFN, which scales a GFlowNet policy's logits by a learned function of the inverse temperature to stabilize and speed up training of temperature-conditional GFlowNets and improve offline generalization and online mode discovery.
Problem
Temperature-conditional GFlowNets, which sample proportional to R(x)^beta, are hard to train stably across temperatures because high-beta targets produce very different gradient profiles and logit magnitudes, and existing conditioning (layer-conditioning) is numerically difficult.
Approach
Add a direct pathway that divides the policy logits by a learned softmax temperature T = f(beta) (logit scaling), with the policy network independent of beta, optionally combined with layer-conditioning; applied with a temperature-conditional trajectory balance objective and an online exploration scheme that samples beta across rounds.
Key results
- Logit-GFN keeps TB loss stable and converges faster than Layer-GFN and unconditional GFN when trained with high beta (beta=5000) on TFBind8 and RNA-binding.
- In online mode seeking, Logit-GFN discovers more modes than GFN, Layer-GFN, MARS, A2C, Soft Q-Learning and PPO across QM9, sEH, TFBind8 and RNA-binding.
- On maximum independent set and maximum clique, Logit-GFN (DB) outperforms GFN (DB), e.g., 19.20 vs 17.93 on MIS (Small).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Scale policy logits directly as a function of temperature | Temperature-conditional GFlowNets · MIS and MC graph combinatorial optimization | best solution size: 17.93 ± 0.69 -> 19.20 ± 0.07; 34.19 ± 0.37 -> 35.65 ± 0.61; 15.56 ± 0.11 -> 16.54 ± 0.20; 29.36 ± 0.01 -> 34.08 ± 3.15 | Helped |
| Add logit scaling to Layer-GFN | Temperature-conditional GFlowNets · Training at high, previously unseen β | TB loss stability | Helped |
| Use Logit-GFN instead of Layer-GFN | Offline generalization · QM9, TFBind8, RNA-binding | generated reward | Helped |
| Use simulated annealing rather than uniform temperature sampling (ablation) | Logit-GFN online exploration · QM9 | number of modes discovered | Helped |
| Add layer-conditioning to Logit-GFN (ablation) | Logit-GFN · Online mode seeking | number of modes discovered | No clear effect |
| Replace uniform exploration temperatures with inverse simulated annealing (ablation) | Logit-GFN online exploration · QM9 | mode discovery efficiency | Mixed |
| Use thermometer encoding for Layer-GFN (ablation) | Layer-GFN · TFBind8, QM9, sEH | number of modes discovered | Mixed |
| Increase the number of gradient steps per batch for unconditional GFlowNets (ablation) | Unconditional GFlowNets · Online mode seeking | number of modes discovered | No clear effect |
| Train Logit-GFN using DB or SubTB instead of TB (ablation) | Logit-GFN · QM9 and TFBind8 | number of modes discovered | Helped |
| Add layer-conditioning to logit scaling (ablation) | Logit-GFN · GridWorld (n=3, H=32), β=2 | L1(×10^-5): 2.777 ± 0.119 -> 2.376 ± 0.231 | Helped |
| Add layer-conditioning to logit scaling (ablation) | Logit-GFN · GridWorld (n=3, H=32), β=3 | L1(×10^-5): 2.130 ± 0.158 -> 1.490 ± 0.075 | Hurt |
Surprising findings
- Layer-conditioning added to logit scaling has only a marginal effect on online mode discovery, and logit-only scaling is competitive with the combined version.
- Thermometer encoding for Layer-GFN helps on TFBind8 but hurts on QM9 and sEH, so its usefulness varies across tasks.
Limitations
- Effectiveness in very high-dimensional or highly complex generative tasks is not fully explored.
- Experiments are limited to biochemical tasks and small graph combinatorial optimization problems.
- Thermometer encoding and layer-conditioning benefits are task-dependent and inconsistent.
Experimental setup
Datasets: QM9 sEH TFBind8 RNA-binding (L14-RNA1) Gridworld maximum independent set maximum clique
Benchmarks: QM9 sEH TFBind8 RNA-binding Gridworld MIS MC
Metrics: training loss (TB) stability 25th percentile and median reward in offline generalization number of modes discovered L1 distance to target distribution
Baselines: GFN (unconditional) Layer-GFN MARS A2C with entropy regularization Soft Q-Learning PPO
Builds on: GFlowNets trajectory balance layer-conditioning prioritized replay training MARS
Closest papers by meaning
- Pre-Training and Fine-Tuning Generative Flow Networks (ICLR 2024 · similarity 0.88)
- Learning Energy Decompositions for Partial Inference in GFlowNets (ICLR 2024 · similarity 0.87)
- On Divergence Measures for Training GFlowNets (NeurIPS 2024 · similarity 0.87)
- Order-Preserving GFlowNets (ICLR 2024 · similarity 0.87)
- gfnx: Fast and Scalable Library for Generative Flow Networks in JAX (NeurIPS 2026 · similarity 0.86)
- To Cool or not to Cool? Temperature Network Meets Large Foundation Models via DRO (ICML 2024 · 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.