liveBook · Manning

Gumbel-softmax

As explained in Graph Neural Networks in Action

Gumbel-Softmax is a differentiable approximation that samples from categorical distributions while enabling gradient-based optimization.

Gumbel-Softmax makes categorical edge predictions usable in gradient-based models. It relaxes encoder logits into a continuous, differentiable distribution during training, then supports discrete edge samples during testing. section 6.5.4

In neural relational inference, the mechanism follows this pipeline:

  • The encoder produces logits representing possible edge types, such as category 0 for disconnected and category 1 for connected. section 6.4 section 6.5.4
  • Gumbel-Softmax converts those logits into edge distributions that remain differentiable, allowing reconstruction and KL-divergence losses to backpropagate. section 6.4 section 6.5.4
  • The temperature controls distribution sharpness; the implementation uses 0.5, while hard sampling is enabled for discrete test-time edges. section 6.5.4
  • The decoder receives sampled edge representations to reconstruct node data or predict future sensor trajectories. section 6.4

PyTorch NRI sampling

Flatten the encoder logits by edge type, call F.gumbel_softmax(logits, tau=0.5, hard=hard_sample), and reshape the result before passing it to the decoder. Training keeps samples differentiable; evaluation uses hard=True.

How Graph Neural Networks in Action covers it

How deep it goes
Explains categorical relaxation, temperature, hard sampling, and PyTorch integration without deriving the Gumbel distribution mathematically.
Examples
Uses binary connected-or-disconnected edges, sensor trajectories, encoder logits, decoder inputs, and future-prediction code.
What it assumes
Assumes encoder logits and categorical edge types while clarifying why discrete graph predictions require a continuous relaxation for optimization.
What you can do after
Apply via PyTorch's F.gumbel_softmax to train differentiable edge predictions and produce hard graph samples for evaluation.
How it connects
Connects categorical distributions, adjacency matrices, variational autoencoders, neural relational inference, and backpropagation.
Cover of Graph Neural Networks in Action
liveBookmanning.com
Read Graph Neural Networks in Action on liveBook Get the book at manning.com

Related concepts

This explainer was generated by Manning's knowledge graph from the text of Graph Neural Networks in Action. The section links point to the parts of the book it draws on.