Reflections on ICML 2026 in Seoul

Calibrating Generative Models to Distributional Constraints

Henry D. Smith, Nathaniel L. Diamant, Brian L. Trippe · Paper · Code

In this paper, the focus is on distributional miscalibration in generative models. A model may produce plausible individual outputs while generating inaccurate or undesirable aggregate patterns, such as underrepresenting certain protein structures, overproducing particular image classes, or associating professions disproportionately with one gender. They formulate calibration as the constrained optimisation problem

$$p_{{\theta }^{\ast }}=\underset{p_{\theta }}{arg\hspace{0.17em}min}{D}_{KL}(p_{\theta }\parallel p_{{\theta }_{\mathrm{base}}})\text{subject to}{𝔼}_{p_{\theta }}[h(x)]=h^{\ast }.$$

The objective is therefore to satisfy specified expectations over output features while remaining as close as possible to the original model. This framework can be understood as a form of regularised moment matching, in which selected expectations of the generated distribution are aligned with target values while KL divergence limits deviation from the base model. The formulation is particularly useful because the constraints can represent class frequencies, scientific measurements, or other population-level properties, and the functions used to evaluate them need not be differentiable.

Because directly solving the constrained problem is generally intractable, the paper introduces two approximate methods called Calibrating Generative Models, or CGM. CGM-relax replaces exact constraint satisfaction with a penalty for miscalibration and uses KL regularisation to prevent excessive deviation from the original model. Its regularisation parameter controls the trade-off between calibration and fidelity. CGM-reward instead uses maximum-entropy theory to construct a reward function whose optimum should exhibit the desired aggregate statistics, allowing calibration to be treated as KL-regularised reward fine-tuning. This establishes a useful connection with exponential tilting and minimum-relative-entropy methods, which adjust a distribution as little as possible while incorporating new constraints. The theoretical connection between distributional constraints, maximum entropy, and reward fine-tuning is one of the paper's strongest contributions. (Jaynes's classical maximum entropy principle says “Among all probability distributions satisfying the specified constraints, choose the one with the largest Shannon entropy.”) The authors also derive unbiased Monte Carlo gradient estimators, making CGM applicable to several model classes as long as they support sampling and likelihood-gradient evaluation.

Across protein design, image generation, and language modelling, CGM-relax substantially improves distributional calibration. It makes generated protein structures more representative of natural proteins, balances animal classes in TarFlow, and reduces gender imbalance in Gemma-2-9B-IT by more than threefold relative to the prior fine-tuning baseline. These improvements involve some trade-offs, including increased protein-design failures, occasional visual artefacts, and inconsistent transfer to professions excluded from language-model fine-tuning. Nevertheless, CGM-relax is generally more reliable and scalable than CGM-reward, reducing calibration error across as many as 1,000 synthetic constraints. CGM-reward performs well in lower-dimensional settings, though estimating its normalising constant requires large sample counts, which forced the protein experiments to use a substantially reduced constraint set. Although neither method guarantees exact calibration and both find extremely rare outputs challenging, the experiments demonstrate that substantial changes in aggregate behaviour can be achieved while largely preserving the capabilities of the original models.

Overall, the paper provides a clear, theoretically grounded, and broadly applicable approach to distributional calibration. Its central contribution is to shift attention from rewarding desirable individual samples to controlling properties of the sampling distribution as a whole. The combination of a principled optimisation objective, practical algorithms, and experiments across several domains makes the work persuasive. While further research is needed to improve rare-event calibration and extend the framework to models without tractable likelihoods, these limitations create promising directions for future work. CGM represents an important step towards generative models whose population-level behaviour can be deliberately controlled without requiring them to be retrained from scratch.

Rotary Position Encodings for Graphs

Isaac Reid, Arijit Sehanobish, Cederik Höfs, Bruno Mlodozeniec, Leonhard Vulpius, Federico Barbero, Adrian Weller, Krzysztof Choromanski, Richard E. Turner, Petar Veličković · Paper

The authors of this spotlight paper investigate whether rotary position encoding (RoPE), originally developed for sequential data and widely adopted in language and vision transformers, can be extended to graph-structured inputs. Positional encoding is more difficult for graphs because their nodes have no canonical ordering or coordinate system. A node's significance instead depends on its connections and its position within the wider topology. Existing graph transformers commonly address this using Laplacian eigenvectors as absolute positional features or pairwise quantities such as shortest-path distance as relative attention biases. Although effective, pairwise biases generally require constructing an $N\times N$ attention matrix and are therefore incompatible with linear-attention methods. The paper addresses this limitation by introducing Wave-Induced Rotary Encodings, or WIRE, which incorporates graph structure directly into the query and key representations.

WIRE begins with the graph Laplacian $L=D-A$, where $A$ is the adjacency matrix and $D$ the diagonal degree matrix, and its eigendecomposition $L=U\Lambda U^{\top }$. The columns of $U$ are the Laplacian's eigenvectors $u_0,\dots ,u_{N-1}$, each oscillating across the graph at the frequency given by its eigenvalue, so that $u_k[i]$ is the value of the $k$th eigenvector at node $i$. The lowest $m$ of them provide spectral coordinates that vary smoothly across the graph, allowing each node to be represented as $r_i={[u_k[i]]}_{k=0}^{m-1}$. These coordinates supply the angles for rotations applied to two-dimensional blocks of the node's query and key vectors:

$$\begin{array}{c}\mathrm{RoPE}(r_i)z_i≔\underset{n=1}{\overset{d/2}{⨁}}\rho ({\theta }_n){[z_i]}_{2n-2:2n-1},\\ \rho (\theta )≔\left(\begin{array}{cc}\cos \theta & -\sin \theta \\ \sin \theta & \cos \theta \end{array}\right),{\theta }_n≔{\omega }_n^{\top }{r}_i,z_i\in \{q_i,k_i\}.\end{array}$$

Here ${[z_i]}_{2n-2:2n-1}$ is the $n$th two-element block of the token, $⨁$ denotes the direct product of those per-block rotations, and the frequencies ${\omega }_n\in {ℝ}^m$ are learnable or fixed — they are WIRE's only learnable parameters, at $dm/2$ per layer. WIRE is therefore exactly RoPE, fed spectral coordinates in place of sequence or Cartesian positions. As with ordinary RoPE, the resulting query–key interaction depends on the difference between the nodes' coordinates, $r_j-r_i$, rather than their absolute positions. This gives attention a useful structural inductive bias without adding an explicit pairwise bias matrix. The method is permutation equivariant with respect to node ordering—up to sign flips and rotations within degenerate eigen-subspaces—recovers conventional RoPE as a special case on grid graphs, and, under particular assumptions, causes attention to depend asymptotically on effective resistance, a graph distance that reflects the number and strength of paths connecting two nodes. These theoretical connections provide a convincing explanation for why spectral rotations should help attention recognise graph structure.

The experiments show that WIRE is particularly effective when predictions depend strongly on graph topology. On a synthetic shortest-path task, it reduces test error from 0.065 to 0.038, while experiments involving monochromatic connected subgraphs show that higher-frequency coordinates become increasingly useful as graph topology grows more complex. WIRE also improves transformer performance on point-cloud classification and segmentation and produces consistent gains across a broad collection of standard graph benchmarks. Importantly, it works with Performer-style linear attention and sometimes closes much of the performance gap between linear and full quadratic attention. The additional learned parameters are small relative to the overall model, making the method an economical architectural modification. Nevertheless, the gains on some conventional benchmarks are modest, and WIRE does not always outperform full-attention models. Computing Laplacian eigenvectors introduces a preprocessing cost, though the authors argue this is largely mitigated in practice: approximate methods such as Lanczos iteration suffice when only the lowest few eigenvectors are needed, and the cost is amortised by the structural features practitioners typically compute anyway. Spectral coordinates are also subject to eigenvector ambiguity, although the paper notes that this is a well-studied problem rather than one specific to WIRE, and reports that standard remedies such as SignNet make little difference to its results. The point-cloud experiments are reported using a single seed, so further replication would strengthen the empirical evidence.

Overall, the paper provides an elegant and practical extension of rotary position encoding to graphs. Its main strength is that it combines ideas from spectral graph theory with the computational advantages of RoPE, allowing structural information to be incorporated without sacrificing compatibility with scalable attention mechanisms. The theoretical results connect the method to familiar notions of position on grids and distance on graphs, while the experiments demonstrate its usefulness across synthetic, geometric, and real-world tasks. Although future work could examine spectral stability, directed or dynamic graphs, and larger-scale applications, WIRE offers a lightweight and well-motivated foundation for scalable graph transformers. More broadly, the paper shows how successful transformer techniques can be adapted to new data structures by identifying the appropriate mathematical notion of position rather than assuming that all inputs possess a natural sequence.

Multivariate Distributional Reinforcement Learning Using Sliced Divergences

Baptiste Debes, Tinne Tuytelaars · Paper · Code

The authors of this paper address an important challenge in reinforcement learning: how to model the full range of possible future outcomes when rewards have several dimensions. Whereas conventional reinforcement learning usually predicts an average return, distributional reinforcement learning represents the entire return distribution. This offers a richer understanding of uncertainty, correlations and trade-offs, but comparing multivariate distributions can quickly become computationally demanding.

The paper introduces Sliced Distributional Reinforcement Learning, or SDRL, as an elegant solution. Instead of comparing complex multivariate distributions directly, SDRL projects them onto a collection of one-dimensional directions. Each pair of projected distributions can then be compared using an efficient, well-understood measure such as the Wasserstein distance, the Cramér distance—the squared $L^2$ distance between cumulative distribution functions, and a close sibling of Wasserstein-1, which is the corresponding $L^1$ distance—or Maximum Mean Discrepancy. Combining the results across projections provides a practical measure of similarity between the original multivariate distributions:

$$\begin{array}{c}{S\Delta }_p^p(\mu ,\nu )={\int }_{S^{d-1}}{\Delta }^p({(P_{\theta })}_{\#}\mu ,{(P_{\theta })}_{\#}\nu )d\sigma (\theta ),\\ \mathrm{MS}\Delta (\mu ,\nu )=\underset{\theta \in S^{d-1}}{sup}\Delta ({(P_{\theta })}_{\#}\mu ,{(P_{\theta })}_{\#}\nu ).\end{array}$$

The mechanism is easier to picture than to read. A direction $\theta$ is a unit vector, and projecting all of a distribution's mass onto it leaves a one-dimensional shadow; comparing two multivariate distributions then reduces to comparing their shadows, which is cheap because one-dimensional divergences have efficient estimators—in one dimension the Wasserstein distance is essentially a sort. Nothing is lost in principle, because two distributions agree exactly when all of their one-dimensional projections agree. Formally, $P_{\theta }(x)=⟨\theta ,x⟩$ is that projection and ${(P_{\theta })}_{\#}\mu$ is the resulting one-dimensional distribution, while $\Delta$ is any base divergence with a tractable estimator in one dimension. The two appearances of the same letter are unrelated: $S^{d-1}$ is the unit sphere of projection directions, with $\sigma$ the uniform measure on it, whereas the leading $S$ marks the sliced lift of the base divergence and $\mathrm{MS}$ its max-sliced counterpart. The two lines then differ only in how they combine directions: the first averages over the sphere, estimated in practice by Monte Carlo over a finite set of draws whose projected subproblems are independent and so can be evaluated in parallel; the second keeps only the single direction along which the two distributions disagree most.

A major strength of the work is that it connects this practical method with a substantial theoretical foundation. The authors show that uniform slicing retains the convergence properties required for reliable distributional learning when all reward dimensions share a scalar discount factor. They also introduce a max-sliced variant that searches for the direction in which two distributions differ most. This provides valuable contraction guarantees—the assurance that repeated Bellman updates provably converge to a single fixed point, rather than depending on where they started—for more general settings involving dense discount matrices, where the discount can weight and mix reward dimensions rather than scaling them all alike. The two properties do not, however, come together: max-slicing forfeits the unbiased sample-gradient property that uniform slicing preserves, and whether any divergence can deliver both is left as a central open question raised by the paper.

The paper also offers a particularly useful analysis of how these objectives behave during training. It demonstrates that theoretical convergence at the population level does not automatically guarantee reliable learning from individual sampled transitions. This distinction helps explain why some mathematically appealing objectives can struggle in ordinary temporal-difference learning. Uniformly sliced Cramér and MMD objectives behave especially well because their sample-based gradients remain aligned with the underlying population objective.

Experiments on a stochastic chain, image-based maze environments and Atari games support the theoretical analysis. Sliced Cramér and sliced MMD accurately reproduce multivariate return distributions, while the Atari results show that sliced objectives can also support strong control performance. The experiments additionally suggest that the method can work effectively with relatively few projection directions, making it attractive from a computational perspective. One result cuts against the theory: sliced Wasserstein-2 performs remarkably well on Atari despite its biased gradients and poor distributional matching elsewhere, which suggests that accurate distributional evaluation and strong control performance need not go together.

Overall, the paper presents SDRL as a principled and practical framework for extending distributional reinforcement learning to multivariate rewards. Its most encouraging practical conclusion is that sliced Cramér provides a strong default: it is efficient, compatible with standard sample-based training and supported by meaningful convergence guarantees. More broadly, the work creates a flexible foundation on which future multivariate distributional learning methods can be developed.

Magnitude Distance: A Geometric Measure of Dataset Similarity

Sahel Torkamani, Henry Gouk, Rik Sarkar · Paper · Code

Here, “magnitude distance” is introduced as a promising new way to compare finite datasets. The work addresses an important problem in machine learning: determining whether two datasets have the same underlying structure, particularly when they contain many dimensions. Existing approaches such as Maximum Mean Discrepancy and Wasserstein distance are useful, but they can become less discriminative in high-dimensional spaces or overly sensitive to outliers. The first of these failures is a form of concentration of measure: as dimensionality grows, pairwise distances between samples become increasingly uniform, and distances built on them lose the contrast needed to separate one distribution from another. The paper locates the effect in the spectrum, observing that Gaussian and exponential kernel matrices acquire low effective rank in high dimensions, so that MMD is dominated by a few leading eigen-directions and in effect compares the samples within a low-dimensional subspace. The authors respond to these challenges by developing a measure based on the mathematical concept of metric-space magnitude, which inverts the similarity matrix and therefore weights directions by inverse eigenvalue, aggregating the part of the spectrum that MMD discards rather than collapsing onto its dominant directions.

Magnitude can be understood as the effective size of a dataset after accounting for its geometry. Points that are close together contribute less independent information, while well-separated points contribute more. The authors construct a similarity matrix from the pairwise distances between observations and use its inverse to calculate this effective size:

$$\begin{array}{c}{\mathrm{Mag}}_X(t)=1^{\top }{K}_t^{-1}1,{(K_t)}_{ij}=e^{-t\Vert x_i-x_j\Vert },\\ d_{\mathrm{Mag}}(X,Y;t)=2{\mathrm{Mag}}_{X\cup Y}(t)-{\mathrm{Mag}}_X(t)-{\mathrm{Mag}}_Y(t).\end{array}$$

Here $K_t$ is the similarity matrix of the dataset at scale $t$, whose entries fall off exponentially with the distance between points, and $1$ is the vector of ones. The magnitude of a set is thus the sum of the weights solving $K_{t}w=1$, and the distance is the excess magnitude of the union over its parts. The magnitude distance between two datasets therefore measures how much additional geometric information becomes available when the datasets are combined. If their structures are similar, combining them adds relatively little; if they occupy meaningfully different regions or exhibit different patterns, the increase is larger. Readers who know information theory will recognise the form: this is the variation of information with magnitude substituted for entropy and set union for the joint distribution. The analogy has a limit, however, and an instructive one. Entropy is submodular whereas magnitude is not, which is exactly why the variation of information satisfies the triangle inequality and magnitude distance does not.

One of the most attractive features of the method is its scale parameter, $t$. At small values of $t$, the comparison emphasises broad, global structure. As $t$ increases, it becomes more sensitive to local clusters and fine-grained differences. This gives magnitude distance a multiscale interpretation that is especially valuable in high-dimensional analysis. Rather than producing only a single similarity score, it can reveal the geometric resolution at which two datasets begin to differ. That flexibility cuts both ways, however: $t$ has to be swept rather than inferred, the paper offers guidance rather than an automatic rule for choosing it, and a poorly chosen scale degrades the measure rather than merely blunting it.

The theoretical analysis establishes several useful properties. Magnitude distance is symmetric and non-negative, and it behaves predictably as the scale changes. As $t$ approaches zero, the distance also approaches zero because only the coarsest structure is visible. At very large scales, individual points become distinguishable and the measure approaches the size of the symmetric difference between the sets. The method is also insensitive to exact duplicates, reflecting its emphasis on geometric diversity rather than simple observation counts.

The paper's high-dimensional experiments are particularly encouraging. Using samples from Gaussian distributions in spaces of up to 1,000 dimensions, the authors show that Maximum Mean Discrepancy can collapse towards zero as dimensionality increases. Magnitude distance remains informative when the scale is appropriately adjusted, particularly when $t$ is chosen on the order of $1/\sqrt{D}$, where $D$ is the dimensionality of the data. It also displays relatively low variability across repeated samples compared with MMD and sliced Wasserstein distance. These results suggest that the method could provide a stable signal in settings where conventional distances struggle with concentration effects.

Magnitude distance also demonstrates useful robustness to outliers. Wasserstein distance can increase substantially when even a small number of observations are placed far from the rest of the data. In the paper's contamination experiments, magnitude distance remains comparatively stable as the outliers are moved progressively farther away. The authors supplement this empirical result with a theoretical boundedness result under suitable conditions on the magnitude weights. This combination of geometric sensitivity and resistance to extreme observations is one of the paper's strongest contributions.

To illustrate its practical potential, the authors introduce the Magnitude Generative Network, or MagGN. This model uses magnitude distance as a training objective and increases the scale parameter progressively during training. The generator therefore learns coarse structure first before being asked to reproduce finer details, creating a natural geometric curriculum. On MNIST, MagGN produces images that are visually competitive with a WGAN-GP baseline while substantially reducing training time in the reported experiments.

The experiments on more complex images also identify productive directions for further development. On CIFAR-10 cat images, magnitude-based pretraining improves the subsequent training of a WGAN-GP, reducing total training time and improving the reported Inception Score. For CelebA, combining magnitude distances in pretrained feature space and pixel space improves the FID score from 32.34 to 24.32. These results indicate that magnitude distance may be especially valuable when used alongside learned representations or as an efficient pretraining signal.

There are still areas requiring further investigation. Scale selection remains important, exact computation requires a potentially expensive kernel-matrix inversion, and the distance does not satisfy the triangle inequality in dimensions greater than one (so it is not, strictly speaking, a metric at all, notwithstanding the paper's own description of it as a distance metric). Its identity property is also defined through “magnitude equivalence,” meaning that distinct sets can occasionally receive zero distance if their non-zero magnitude weights have the same support. The current empirical evidence is primarily proof-of-concept, so evaluation across larger and more varied real-world datasets would strengthen the practical claims.

Overall, the paper presents an original and mathematically well-motivated approach to dataset comparison. Its multiscale behaviour, high-dimensional discriminability and robustness to outliers make it a compelling complement to MMD and Wasserstein-based methods. Although additional work is needed on scalable approximation and automatic scale selection, magnitude distance offers a flexible new perspective: dataset similarity can be understood not only through probability or pairwise transport, but through the amount of geometric information each dataset contributes.

Recurrent Structural Policy Gradient for Partially Observable Mean Field Games

Clarisse Wibault, Johannes Forkel, Sebastian Towers, Tiphaine Wibault, Juan Duque, George Whittle, Andreas Schaab, Yucheng Yang, Chiyuan Wang, Mike Osborne, Benjamin Moll, Jakob Foerster · Paper · Code

The subject matter of this spotlight paper by Clarisse Wibault et al. is mean field games, a framework that models settings in which a large population of agents responds to the aggregate distribution of everyone else rather than to any individual—such as, an investor reacting to a price the population itself determines, drivers to congestion, households to interest rates. The algorithmic difficulty is a familiar trade-off. Model-free multi-agent reinforcement learning treats the dynamics as a black box and pays for it in gradient variance, while exact dynamic-programming methods scale poorly. Hybrid Structural Methods occupy the middle ground: because individual state and action spaces are typically small and the individual transition dynamics are known, one can compute the expected return exactly over individual transitions and sample only the common noise. The catch is that prior Hybrid Structural Methods were limited to memoryless (Markov) policies, preventing them from handling partially observable mean field games that require history-dependent policies.

Extending them is not merely an engineering exercise. Once policies condition on histories, the population distribution must be maintained over the history space rather than the individual state space, and the number of histories branches exponentially with time. The paper's insight is that in many applications the observation is public—every agent sees the same aggregate signal—so memory can be restricted to the history of shared observations rather than each agent's private trajectory. The authors formalise the condition under which this loses nothing:

$$p({\mu }_t,z_t\mid {\tau }_t)=p({\mu }_t,z_t\mid s_t,o_{0:t}).$$

In words: given an agent's current individual state and the public history, its own past trajectory contains no additional information about the aggregate state. An investor's past trades reveal little about the economy once the price history and current portfolio are known. The resulting algorithm, Recurrent Structural Policy Gradient (RSPG), is the first history-aware Hybrid Structural Method. To the authors' credit the assumption is not oversold: they provide a counter-example in which the equivalence fails and present it as a useful approximation rather than a theorem about mean field games in general.

The experiments cover three partially observable environments—Linear Quadratic, Beach Bar and a heterogeneous-agent Macroeconomics model in the Krusell–Smith tradition, where interest rates and wages are set endogenously by the wealth and income distribution. RSPG is compared against its memoryless counterpart and against Independent PPO, Recurrent PPO and M-OMD, using exploitability against a perfect-information evaluation class and measuring cost in wall-clock time, since environment-step counts are not comparable across these method families. The structural methods converge roughly an order of magnitude faster than the reinforcement-learning baselines, and RSPG is the only method consistently competitive across all three environments.

The more interesting result is behavioural. In Beach Bar, history-aware agents learn to move away before the bar's random closure; in the macroeconomic model they learn to spend down wealth towards the end of the episode, pushing wages down and interest rates up. In that environment the memoryless method attains comparable exploitability yet never acquires the anticipation—so the two agree on the equilibrium's value while disagreeing about the behaviour that produces it. For anyone using these models to reason about how a population behaves rather than merely what it is worth, that gap is the point.

The limitations are stated fairly. The public-information assumption excludes genuinely private signals, which is a real restriction in market settings where informational asymmetry is the phenomenon of interest. The authors also note that lower gradient variance does not translate into lower variance in final exploitability across seeds, since different runs may converge to different equilibria, and that these methods still sample the common noise. Alongside the algorithm the paper contributes MFAX, a JAX framework that separates white-box from sample-based access to the dynamics and reports mean-field updates one to three orders of magnitude faster than existing libraries. That combination—a variance-reduction method that survives partial observability, plus tooling that makes the comparison reproducible—is what makes the work more useful than its benchmark numbers alone suggest.