2a:null 2b:["$","title",null,{"children":"Gumbel–Softmax Distribution | Mattias Villani | Observable"}] 31:T54b,md` The [Gumbel–Softmax Distribution](https://en.wikipedia.org/wiki/Dirichlet_distribution) is a multivariate probability distribution defined over the space of positive vectors that sum to one; that is, ${tex`\boldsymbol{x} = \{x_1, \dots, x_K \}`} where ${tex.block` \sum_{k=1}^K x_k = 1 \text{ and } x_k \geq 0 \text{ for all } k = 1, \dots, K. `} This space is also called the ${tex`K - 1`} [simplex](https://en.wikipedia.org/wiki/Simplex). The density of the Gumbel-Softmax distribution is given by ${tex.block` \mathrm{Gumbel}-\mathrm{Softmax}(\boldsymbol{x} \mid \boldsymbol{\pi}, \tau) = \Gamma(K) \tau^{K-1}\Big(\sum_{k=1}^K\frac{\pi_k}{x_k^\tau}\Big)^{-K}\prod_{k=1}^K\frac{\pi_k}{x_k^{\tau+1}} `} where ${tex`\Gamma(\cdot)`} is the Gamma function, ${tex`\boldsymbol{\pi}=(\pi_1,\ldots,\pi_K)`} is a discrete probability distribution and ${tex`\tau>0`} is the temperature or concentration parameter. A distribution on the unit simplex is typically visualized for the case where ${tex`K=3`} because it's still possible to render it in 2 dimensions using a [ternary plot](https://en.wikipedia.org/wiki/Ternary_plot). Each slider below controls one element of the parameter vector ${tex`\boldsymbol{\pi}`} and the temperature parameter ${tex`\tau>0`}. The ternary plot visualizes the Gumbel-Softmax probability density over the simplex. `32:Td64,ternaryDensity = (density, resolution, options) => { options = Object.assign({ size: 400, margin: { left: 30, top: 30, right: 30, bottom: 30 }, colorScheme: d3.interpolateBlues }, options); const fillScale = d3.scaleSequential(options.colorScheme).domain(d3.extent(density.map(d => d.density))) const svg = d3.create("svg") .attr("width", options.size) .attr("height", options.size); const x = d3.scaleLinear() .domain([0, 1]) .range([options.margin.left, options.size - options.margin.right]) const y = d3.scaleLinear() .domain([0, 1]) .range([options.size - options.margin.bottom, options.margin.top]); const axisBottom = g => g.call(d3.axisBottom(x).ticks(4)) const axisTop = g => g.call(d3.axisTop(x).ticks(4)); const hexagonSize = ((x(1 / (2 * resolution)) - x(0))) / Math.cos(Math.PI / 6) + 1; const corners = [{ x: 0.5, y: Math.sqrt(3)/2 }, { x: 0, y: 0 }, {x: 1, y: 0 }]; const corners_expanded = [{ x: 0.5, y: Math.sqrt(3)/2 + 0.1 }, { x: -0.1, y: -0.05 }, {x: 1.1, y: -0.05 }]; const line = d3.line().x(d => x(d.x)).y(d => y(d.y)); svg.append('defs') .append('clipPath') .attr('id', 'triangle') .append('path') .attr('x', options.size / 2) .attr('y', options.size / 2) .attr('d', line(corners)) svg.selectAll('.border') .data([[corners[0], corners[1]], [corners[1], corners[2]], [corners[2], corners[0]]]) .enter() .append('line') .attr('x1', (d) => x(d[0].x)) .attr('x2', (d) => x(d[1].x)) .attr('y1', (d) => y(d[0].y)) .attr('y2', (d) => y(d[1].y)) .attr('class', 'border') .attr('stroke', 'black'); const hexbin = d3.hexbin(); svg.append('g') .attr('clip-path', 'url(#triangle)') .selectAll(".point") .data(density) .join('path') .attr('transform', d => `translate(${x(corners[0].x * d.x[0] + corners[1].x * d.x[1] + corners[2].x * d.x[2])}, ${y(corners[0].y * d.x[0] + corners[1].y * d.x[1] + corners[2].y * d.x[2])})`) .attr('fill', d => fillScale(d.density)) .attr('d', hexbin.hexagon(hexagonSize)); svg.selectAll('text') .data([[corners_expanded[0], corners_expanded[1]], [corners_expanded[1], corners_expanded[2]], [corners_expanded[2], corners_expanded[0]]]) .join("text") .attr("text-anchor", "middle") .attr("alignment-baseline", "middle") .attr("font-style", "italic") .attr("x", d => (x(d[0].x) + x(d[1].x)) / 2) .attr("y", d => (y(d[0].y) + y(d[1].y)) / 2) .text((d, i) => { if(i == 0) return "x₂"; if(i == 1) return "x₃"; if(i == 2) return "x₁"; }) const removeLine = g => g.select(".domain").remove(); svg.append("g") .attr("transform", `translate(0, ${options.size - options.margin.bottom})`) .call(axisBottom) .call(removeLine) svg.append("g") .attr("transform", `translate(${x(1) + options.margin.right / 2 }, ${y(0) + 26}) rotate(-120)`) .call(axisBottom) .call(removeLine) .call(g => g.selectAll("text").attr("transform", "translate(11, 22) rotate(120)")); svg.append("g") .attr("transform", `translate(${x(0.5) + options.margin.left / 2}, ${y(Math.sqrt(3) / 2) - 26}) rotate(120)`) .call(axisBottom) .call(removeLine) .call(g => g.selectAll("text").attr("transform", "translate(-11, 22) rotate(-120)")); return svg.node(); }33:T624,marginalDensities = (marginalDensity, index, options) => { options = Object.assign({ height: 300, width: 500, margin: { left: 25, top: 10, right: 50, bottom: 40 } }, options); const svg = d3.create("svg") .attr("width", options.width) .attr("height", options.height); const x = d3.scaleLinear() .domain([0, 1]) .range([options.margin.left, options.width - options.margin.right]) const y = d3.scaleLinear() .domain([0, d3.max(marginalDensity.flat().map(d => d.density))]) .range([options.height - options.margin.bottom, options.margin.top]); const line = d3.line() .x(d => x(d.x)) .y(d => y(d.density)); const paths = svg.selectAll(".path") .data([marginalDensity[index]]) .join("path") .attr("class", "path") .attr("d", d => line(d)) .attr("stroke", "steelblue") .attr("stroke-width", 2) .attr("fill", "transparent"); const axisBottom = g => g.call(d3.axisBottom(x).ticks(5)) const axisLeft = g => g.call(d3.axisLeft(y).ticks(5)); svg.append("g") .attr("transform", `translate(0, ${options.height - options.margin.bottom})`) .call(axisBottom); svg.append("g") .attr("transform", `translate(${options.margin.left}, 0)`) .call(axisLeft); svg.append("text") .attr("transform", `translate(${options.margin.left + (options.width - options.margin.left - options.margin.right) / 2}, ${options.height - 5})`) .style("font-style", "italic") .text(index == 0 ? "x₁" : index == 1 ? "x₂" : "x₃") return svg.node(); }34:[] 2e:{"id":"413917f9c497edac","slug":"gumbel-softmax-distribution","title":"Gumbel–Softmax Distribution","thumbnail":"c501ba4c042550c858da9d6a39dba6327d0039a75466334ba494586ed8b11938","access_level":"public","create_time":"$D2026-01-13T21:16:15.143Z","update_time":"$D2026-01-14T14:58:32.708Z","view_time":null,"first_public_time":"$D2026-01-13T21:46:45.474Z","fork":{"id":"3b61b39228e56c0c","slug":"dirichlet-distribution","title":"Dirichlet Distribution","owner":{"id":"7592a7ab121f74a0","login":"mattiasvillani","name":"Mattias Villani"}},"stars":0,"starred":false,"owner":{"type":"team","id":"7592a7ab121f74a0","login":"mattiasvillani","name":"Mattias Villani","avatar_url":"https://avatars-next.observableusercontent.com/a440aaf39a3abefb86cbcc758fd7a5ab42555ae68d6825b52ec006f385a94d1d","tier":"pro_2024"},"version":790,"publish_version":null,"role":null,"authored":false,"authors":[{"type":"individual","id":"56499be8cf43f0ef","login":"mattiasvillani","name":"Mattias Villani","avatar_url":"https://avatars-next.observableusercontent.com/a440aaf39a3abefb86cbcc758fd7a5ab42555ae68d6825b52ec006f385a94d1d","tier":null}],"members":[{"type":"individual","id":"56499be8cf43f0ef","login":"mattiasvillani","name":"Mattias Villani","avatar_url":"https://avatars-next.observableusercontent.com/a440aaf39a3abefb86cbcc758fd7a5ab42555ae68d6825b52ec006f385a94d1d","tier":null,"role":"editor"},{"type":"team","id":"7592a7ab121f74a0","login":"mattiasvillani","name":"Mattias Villani","avatar_url":"https://avatars-next.observableusercontent.com/a440aaf39a3abefb86cbcc758fd7a5ab42555ae68d6825b52ec006f385a94d1d","tier":"pro_2024","role":"viewer"}],"comments":[],"custom_thumbnail":false,"trashed":false,"subscribed":null,"forks":0,"collection":{"id":"40c8ef5af1a51aff","slug":"distributions","title":"Statistical distributions","owner":{"id":"7592a7ab121f74a0","login":"mattiasvillani","name":"Mattias Villani"}},"type":"notebook","body":{"title":"Gumbel–Softmax Distribution","cells":[{"id":706,"value":"# Gumbel–Softmax Distribution","pinned":false,"hidden":false,"mode":"md"},{"id":0,"value":"$31","pinned":false,"hidden":false,"mode":"ojs"},{"id":11,"value":"viewof inputs = Inputs.form({\n pi1: Inputs.range([0.01, 0.99], { label: tex`\\pi_1`, value: 0.3, step: 0.01 }),\n pi2: Inputs.range([0.01, 0.99], { label: tex`\\pi_2`, value: 0.5, step: 0.01 }),\n tau: Inputs.range([0.01, 5], {value: 1, step: 0.01, label: tex`\\tau`})\n})","pinned":false,"hidden":false,"mode":"ojs"},{"id":491,"value":"{\n if (inputs.pi1 + inputs.pi2 <= 1){\n return ternaryDensity(density, densityResolution, { size: 400, colorScheme: d3.interpolateReds})\n}else{\n return md`probabilities outside the unit simplex`\n}\n}","pinned":false,"hidden":false,"mode":"ojs"},{"id":774,"value":"{\nif (inputs.pi1 + inputs.pi2 <= 1){\n return legend({color: d3.scaleSequential(d3.extent(density.map(d => d.density)), d3.interpolateReds), title: \"Density\"})\n}else{\n return md``\n}\n}","pinned":false,"hidden":false,"mode":"ojs"},{"id":716,"value":"
\n
\n
\n
\n
\n
\n
\n
\n
\n
\n
\n
\n
\n
\n
\n
\n\n---\n## 🙈 Appendix 🙈\n\nJust some code and details that you probably want to ignore.","pinned":false,"hidden":false,"mode":"md"},{"id":15,"value":"pi3 = 1 - inputs.pi1 - inputs.pi2","pinned":false,"hidden":false,"mode":"ojs"},{"id":485,"value":"$32","pinned":false,"hidden":false,"mode":"ojs"},{"id":553,"value":"$33","pinned":false,"hidden":false,"mode":"ojs"},{"id":19,"value":"gumbelsm = (x, Pi, tau) => {\n const K = Pi.length\n let produkt = 1;\n let summa = 0;\n for(let i = 0; i < x.length; i++){\n produkt *= Pi[i] / Math.pow(x[i], tau + 1);\n summa += Pi[i] / Math.pow(x[i], tau);\n }\n return mathfn.gamma(K)*Math.pow(tau, K-1) * Math.pow(summa,-K) * produkt;\n}","pinned":false,"hidden":false,"mode":"ojs"},{"id":533,"value":"density = {\n const density = [];\n for(let i = 1; i < densityResolution; i++) {\n for(let j = 1; j < densityResolution - i; j++) {\n const x1 = i / densityResolution ;\n const x2 = j / densityResolution;\n const x3 = Math.max(0, 1 - x1 - x2);\n const x = [ x1, x2, x3 ];\n\n density.push({ x, density: gumbelsm(x, Pi, inputs.tau) });\n }\n }\n return density;\n}","pinned":false,"hidden":false,"mode":"ojs"},{"id":535,"value":"densityResolution = 40","pinned":false,"hidden":false,"mode":"ojs"},{"id":563,"value":"betaDensity = (x, alpha, beta) => (Math.pow(x, alpha - 1) * Math.pow(1 - x, beta - 1)) / mathfn.beta(alpha, beta)","pinned":false,"hidden":false,"mode":"ojs"},{"id":63,"value":"Pi = [inputs.pi1, inputs.pi2, pi3]","pinned":false,"hidden":false,"mode":"ojs"},{"id":31,"value":"beta = (alpha) => {\n let result = 1;\n for(let i = 0; i < alpha.length; i++) {\n result *= mathfn.gamma(alpha[i]);\n }\n result = result / mathfn.gamma(d3.sum(alpha));\n return result;\n}","pinned":false,"hidden":false,"mode":"ojs"},{"id":51,"value":"margin = ({ left: 30, top: 30, right: 30, bottom: 30 })","pinned":false,"hidden":false,"mode":"ojs"},{"id":8,"value":"height = 400","pinned":false,"hidden":false,"mode":"ojs"},{"id":29,"value":"mathfn = require('https://bundle.run/mathfn@1.1.0')","pinned":false,"hidden":false,"mode":"ojs"},{"id":4,"value":"d3 = require(\"d3@v6\", \"d3-hexbin@0.2\")","pinned":false,"hidden":false,"mode":"ojs"},{"id":179,"value":"import {Range} from \"@observablehq/inputs\"","pinned":false,"hidden":false,"mode":"ojs"},{"id":275,"value":"import {legend} from \"@d3/color-legend\"","pinned":false,"hidden":false,"mode":"ojs"}],"files":"$W34","stdlib":"1","titling":"auto"}} 2f:"$2e:body"