Daniel Bashir

Wasserstein gradient flows meet policy optimization

Sun, 12 Oct 2025 00:00:00 +0000

I’ve really enjoyed two papers recently: “Policy Optimization as Wasserstein Gradient Flows” (Zhang et al. 2018) and “Wasserstein Policy Optimization” (Pfau et al. 2025). Collectively, they present an interesting perspective on the learning of optimal policies in RL.

Following Gradient Fields in Action Space

I implemented Wasserstein Policy Optimization (WPO), a reinforcement learning algorithm that replaces standard policy gradient’s “move probability toward high rewards” with something more sophisticated: follow the gradient field of the value function in action space.

Standard Policy Gradient and its Problems

The standard policy gradient for a Gaussian policy $\pi(a \mid s) \sim \mathcal{N}(\mu,\sigma^2)$ updates as $\nabla_{\theta}J = \mathbb{E}_{\pi} \left[ Q(s,a) \cdot \nabla_{\theta} \log\pi(a \mid s) \right]$ The Q-value $Q(s,a)$ is a scalar weight, so high Q-value actions amplify their log probability gradient. This is intuitive: make good actions more likely. However, this fails to include some important information—$Q(s,a)$ only tells us that an action is good, not which direction in action space leads to better actions.

Use Action Gradients!

WPO uses $\nabla_a Q(s,a)$ instead: the gradient of $Q$ w/r/t/ the action itself: $\nabla_{\theta}J = \mathbb{E}_{\pi} \left[ (\nabla_a \log \pi) \cdot (\nabla_a Q) \cdot \nabla_{\theta}\mu \right]$ For a Gaussian policy, $\nabla_a \log\pi(a \mid s) = -(a - \mu) / \sigma^2$, giving $\text{WPO signal} = \frac{-(a-\mu)}{\sigma^2} \cdot \frac{\partial Q}{\partial a}$ In words: the policy mean follows the gradient field of $Q$. if $\partial Q / \partial a > 0$ ($Q$ increasing rightward), the update pushes $\mu$ in that direction, scaled by how far the sampled action was from the mean.

The Zero-Variance Property

We can read off a simple diagnostic from the Gaussian policy formulation above: the WPO signal must be exactly zero when $a = \mu$ (since the numerator becomes zero). The paper highlights that WPO exhibits a zero-variance property when the Q-function is locally linear in action space. This is an important theoretical advantage: where standard Policy Gradient updates still struggle with high variance in this common scenario, the WPO update becomes highly consistent, leading to more stable and efficient training.

Deriving WPO from Wasserstein Gradient Flows

Step 1: Define the distance metric

WPO’s update comes from optimal transport theory. We want to do gradient descent on policies, but policies are probability distributions, not points in Euclidean space. So we need a way to measure distance between distributions.

The 2-Wasserstein metric defines distance as: $W_2^2(\pi_0, \pi_1) = \inf_{\gamma \in \Gamma(\pi_0,\pi_1)} \int \vert a-b \vert^2 d\gamma(a,b)$

Think of this as: to morph distribution $\pi_0$ into $\pi_1$, we transport probability mass from actions in $\pi_0$ to actions in $\pi_1$. The coupling $\gamma(a,b)$ says “move mass from action $a$ to action $b$.” The Wasserstein distance is the minimum cost of such a transport plan, where cost = squared distance × amount of mass moved.

Why this metric? Unlike KL divergence (which cares about probability ratios), Wasserstein cares about geometric distance in action space. If two policies differ slightly (mean shifts 0.01), Wasserstein distance is small. This makes gradients smooth and well-behaved.

Step 2: Gradient flow in distribution space

Now we want to do steepest descent on our objective $J[\pi] = \mathbb{E}_\pi[Q(s,a)]$ while measuring “steepness” using Wasserstein distance. The calculus of variations tells us the steepest descent flow satisfies:

\[\frac{\partial\pi}{\partial t} = -\nabla_a \cdot \left( \pi \nabla_a \frac{\delta J}{\delta\pi} \right)\]

Let’s unpack this PDE:

$\frac{\partial\pi}{\partial t}$: How the policy distribution evolves over time
$\frac{\delta J}{\delta\pi}$: Functional derivative of $J$ w/r/t/ $\pi$ (how much would $J$ change if we perturbed $\pi$ at action $a$?)
$\nabla_a$: Gradient in action space, not parameter space
The divergence form ($\nabla_a \cdot$) ensures probability mass is conserved—like the mass of fluid flowing through a pipe. It means that any change in the policy’s probability density ($\frac{\partial\pi}{\partial t}$) at one action must be exactly balanced by a corresponding flow ($\pi\nabla_a…$) of mass into or out of that action.

This is called the continuity equation for probability flow, analogous to fluid dynamics.

Step 3: Plug in the RL objective

For our objective $J[\pi] = \mathbb{E}_\pi[Q(s,a)]$, the functional derivative is just $Q(s,a)$ itself. Think of it this way: if you slightly increase the probability $\pi$ of a single action $a$ by a tiny amount $\delta\pi(a)$, how much does the expected reward $J$ change? Since $J$ is the average reward, the change $\delta J$ is simply that tiny change in probability $\delta\pi(a)$ multiplied by the reward $Q(s,a)$ associated with that action. This is exactly what the functional derivative $\frac{\delta J}{\delta\pi} = Q(s,a)$ captures.

So our flow becomes: $\frac{\partial\pi}{\partial t} = -\nabla_a \cdot (\pi\nabla_a Q)$

Expanding the divergence: $\frac{\partial\pi}{\partial t} = -\nabla_a \pi \cdot \nabla_a Q - \pi \nabla_a^2 Q$

The second term $\nabla_a^2 Q$ (Hessian) is expensive to compute, so WPO drops it (first-order approximation): $\frac{\partial\pi}{\partial t} \approx -\nabla_a \pi \cdot \nabla_a Q$

Step 4: Project onto parametric policies

The PDE above describes evolution in the infinite-dimensional space of all probability distributions. But we parameterize our policy with a neural network $\pi_\theta$. To project the continuous flow onto our finite parameters $\theta$, we use:

\[\frac{d\theta}{dt} = \mathbb{E}_{\pi_\theta}\left[\nabla_\theta \log\pi_\theta(a \mid s) \cdot \nabla_a Q(s,a)\right]\]

This is the WPO update! It translates the ideal flow direction given by the PDE into an update for the neural network parameters $\theta$:

$\nabla_a Q(s,a)$ from the WGF tells us the most rewarding direction in action space to move the policy mass.
$\nabla_{\theta}\log\pi_{\theta}(a \mid s)$ (the score function) tells us how to change the parameters $\theta$ to push probability mass toward action $a$.

Their product gives the parameter update direction, which makes the policy follow the gradient field of the Q-function in action space.

For a Gaussian policy $\pi(a \mid s) = \mathcal{N}(\mu_\theta, \sigma_\theta^2)$, we have $\nabla_a \log\pi = -(a-\mu)/\sigma^2$, recovering our earlier formula.

A notable equivalence

The paper reveals an interesting result: for the simplest case, a single-variate Gaussian policy, the expectd WPO parameter update is equivalent to the expected standard Policy Gradient update (see Section 4.1 in the WPO paper).

This suggests that for basic continuous control, there’s essentially one optimal expected update direction. But WPO still has benefits:

As noted above, WPO achieves the same expected result with lower variance than PG, especially when the Q-function is linear in action space.
The equivalence breaks down when we move to more complex, multi-modal policies like Mixture-of-Gaussians. The WGF foundation ensures WPO’s update is a valid steepest-descent direction for any policy distribution—WPO pushes different components of the MoG towrad different local action optima, allowing the policy to find multiple action-space optima. PG struggles with this and can become unstable.

Implementation Notes

While the theory in the paper is general, the implementation in both Acme and my version uses Gaussians—this is pretty standard, since Gaussian policies are commonly used for continuous control. I did not test this on DeepMind’s suite but a few environments in the gymnasium repo.

I wrote my implementation in PyTorch, which makes the computation of $(\delta_a \log\pi)\cdot(\delta_a Q)$ rather tricky. It requires gradients through two paths:

# sample from gaussian policy
actions = mean + std * torch.randn_like(mean)

# path 1: \delta_a log \pi via vector-Jacobian product
actions_copy = actions.detach().requires_grad_(True)
log_pi = -0.5 * (((actions_copy - mean) / std)**2 + 
                  2*log_std + np.log(2*np.pi)).sum(-1)

log_pron_grad = torch.autograd.grad(
    log_pi.sum(),
    actions_copy,
    create_graph=True # we need gradients w/r/t/ policy params
)[0]

# path 2: \delta_a Q
actions_for_q = actions.detach().requires_grad_(True)
q_values = critic(state, actions_for_q)
q_grad = torch.autograd.grad(q_values.sum(), actions_for_q)[0]

# full WPO loss: -E[(\delta_a log \pi) * (\delta_a Q)]
wpo_loss = -(log_prob_grad * q_grad.detach()).sum(-1).mean()

Recall that we have two networks and two paths here: the actor $\pi_{\theta}$ which determines the policy and its parameters ($\theta$) which is what we want to update; and the critic $Q(s,a)$ which estimates the action’s value and is updated by minimizing a temporal difference error. The WPO update is designed only to update the actor.

The term q_grad ($\nabla_a Q$) relies on the Critic network’s parameters. Without .detach(), PyTorch sees a complete path from the final wpo_loss to the critics weights. .detach() tells PyTorch to treat the value of q_grad as a fixed numerical constant during this backprop step, so it doesn’t flow any gradients backward through this value into the Critic network.

We use create_graph=True because log_prob_grad will be multiplied by q_grad, and we need $\nabla_{\theta}(\text{log_prob_grad} \cdot \text{q_grad})$. Without it, the gradient won’t flow to the policy parameters.

Understanding Vector-Jacobian Products

The computation of $\nabla_a \log\pi(a \mid s)$ deserves more explanation. We need to compute:

\[\nabla_a \log\pi(a \mid s) = \frac{\partial \log\pi(a \mid s)}{\partial a}\]

This is a gradient of a scalar (log-probability) w/r/t/ a vector (action). If action has dimension $d$, this gives us a $d$-dimensional gradient vector.

In JAX, this is done explicitly with a vector-Jacobian product (VJP):

log_probs, log_proj_vjp = jax.vjp(distribution.log_prob, actions)
log_prob_grad = log_prob_vjp(jnp.ones_like(log_probs))[0]

Some more details (I don’t know JAX very well lol)

The Jacobian of $\log\pi: \mathbb{R}^d \to \mathbb{R}$ is just the gradient (since output is scalar): $J = \nabla_a \log\pi \in \mathbb{R}^{1 \times d}$

A vector-Jacobian product computes $v^T J$ for some vector $v$. When $v = 1$ (a scalar), we get: $v^T J = 1 \cdot \nabla_a \log\pi = \nabla_a \log\pi$

So vjp_fn(ones) literally returns the gradient.

Why use VJP instead of just calling .grad()? We need to multiply this gradient by $\nabla_a Q$ and then backpropagate through that product. The VJP function retains the computational graph, allowing us to compute: $\nabla_\theta \left[(\nabla_a \log\pi) \cdot (\nabla_a Q)\right]$

This is second-order differentiation: we’re differentiating through a gradient.

In PyTorch, we achieve the same thing with create_graph=True:

log_prob_grad = torch.autograd.grad(
    log_pi.sum(),      # scalar output
    actions_copy,       # vector input
    create_graph=True   # keep graph for second-order gradients
)[0]

Without create_graph=True, PyTorch would compute the gradient but throw away the graph needed to backprop through it. With it, we can compute $\nabla_\theta[f(\nabla_a \log\pi)]$ for any function $f$.

The JAX and PyTorch approaches are mathematically equivalent—both compute second-order gradients through different APIs.

Experiments

I implemented WPO in PyTorch and tested it on Pendulum-v1. WPO has two key hyperparameters: $\epsilon_{\text{mean}}$ and $\epsilon_{\text{stddev}}$, which bound how much the policy can change per update via KL divergence constraints. The paper uses $\epsilon = 0.01$ as default, but how sensitive is the algorithm to this choice?

I swept epsilon from 0.001 (very tight constraint) to 0.05 (very loose), running 5 random seeds per configuration:

Figure 1: Six-panel analysis of epsilon sensitivity. Top row: performance vs epsilon, stability (coefficient of variation), KL mean vs constraint. Bottom row: KL stddev, dual variables (constraint pressure), policy exploration level.

Epsilon	Success Rate	Avg Reward	Range
0.001	40%	-377	[-622, -130]
0.005	80%	-222	[-518, -127]
0.01	60%	-468	[-1570, -139]
0.02	80%	-210	[-404, -127]
0.05	80%	-222	[-503, -133]

Surprisingly, the tightest constraint ($\epsilon$=0.001) performed worst, with only 40% of runs converging. Instead, moderate values ($\epsilon$=0.02-0.05) all achieved ~80% success rates. This contradicts the intuition that “tighter constraint = more stable learning.”

The paper’s default $\epsilon$=0.01 falls in an unstable regime here. Note: These results use slightly simplified settings:

Buffer: 100K transitions (paper uses 2M)
Update frequency: every step (paper uses every 100 steps)
Network: [256, 256] actor, [512, 512] critic (paper uses [256, 256, 128] and [512, 512, 256])

The epsilon sensitivity might differ with the paper’s full configuration, though the qualitative finding—that a sweet spot exists and tighter isn’t always better—likely persists.

The KL “Violation”

The WPO paper, like many modern policy optimization algorithms, includes a KL divergence constraint (bounded by $\epsilon$) as a stability measure. While theoretically, this $\epsilon$ should enforce a hard bound on how much the policy shifts, that isn’t quite what I saw in practice.

All runs in my experiment violate the $\epsilon$ constraint by orders of magnitude—yet most still succeed. If I am not doing something wrong, this would suggest that in practice, $\epsilon$ acts less as a strict maximum bound and more as a hyperparameter that modulates the sensitivity of the dual penalty $\alpha$, which in turn controls the step size.

Figure 2: KL divergence evolution during training for different epsilon values. Solid lines show mean across seeds, shaded regions show ±1 std. Dashed horizontal lines indicate the epsilon thresholds. All runs massively violate constraints, yet most converge.

The KL divergence starts near zero but grows to 100-300 by the end of training. For $\epsilon$=0.001, this represents a violation by a factor of 100,000. The dual variables (Lagrange multipliers α) try to enforce constraints by growing from ~1.7 to ~10.7, but cannot contain the drift. Yet 68% of runs still converge successfully. Hence, the guess about epsilon being a sensitivity parameter.

More to do

I’d like to validate the numerical correctness of my WPO update vs. the JAX implementation. This should be straightforward—run updates through both implementations with identical inputs and verify outputs. I’d be surprised if everything were numerically perfect. I also don’t know what hardware the DeepMind authors used (I’m using an oldie Intel Mac chip), but surely there are lots of things that might stray. I’m also testing with the paper’s full configuration (2M buffer, update every 100 steps, 3-layer networks) to see if epsilon sensitivity changes under those conditions.

some notes on value alignment

Mon, 29 Sep 2025 00:00:00 +0000

Lewis Mumford:

To perfect and extend the range of machines without perfecting and giving humane direction to the organs of social action and social control is to create dangerous tensions in the structure of society (Technics and Civilization, 367).

I’m a big fan of Iason Gabriel’s work. He brings a much-needed lens to questions about how AI systems will affect our social structure and the distribution of goods and resources, and whether they will be used in ways conducive to social and moral good. If we believe AI systems might become a foundational part of how we make decisions and organize parts of society, the work to understand what the resulting effects will be and how we ought to develop our social contracts and systems of governance in the face of those changes is vital. In particular, Gabriel cares about value alignment.

Gabriel understands value alignment in two separate parts: the technical challenge of trying to align AI systems with human values, and the normative question of what or whose values we try to align AI systems with. In “Artificial Intelligence, Values, and Alignment,” Gabriel considers six potential goals for alignment (it is worth noting that this is focused on one-person-one-agent scenarios): (1) instructions, (2) expressed intentions, (3) revealed preferences, (4) informed preferences or desires, (5) interest or well-being, (6) values. On the final goal, Gabriel makes the interesting point that, in practice, AI would have to be aligned with a set of beliefs about value as opposed to value itself. That values play an important role in social life makes Gabriel confident that alignment with a community’s moral beliefs is a good target.

Across his papers, I see Gabriel’s normative center of gravity as liberal egalitarian. Rawlsian justice is applied to sociotechnical systems: in his essay “Toward a Theory of Justice for Artificial Intelligence,” he argues that AI is now part of the basic structure of society. In Rawls’ thinking, this basic structure encompasses how the major social institutions (the organization of the economy, the nature of the family) cohere in a single system, how they assign basic rights and duties, and how they shape the socially-mediated divison of advantages. In Gabriels’ reckoning, this basic structure is a composite of sociotechnical systems. Because AI shapes the fundamental institutions that compose the basic structure of society, they are subject to egalitarian norms of justice.

Another theme in Gabriel’s papers is the move towards identifying a fair, broadly endorsed way to develop principles for the alignment or governance of AI systems. With the recognition that there often are reasonable disagreements between different communities, individuals, or value systems—disagreements tht cannot be resolved by appealing to shared first principles—approaches that can be recognized by all as procedurally fair (even if not all parties are equally satisfied by the outcomes of those procedures) seem just in an important sense.

While Gabriel’s focus on procedures for identifying fair principles feels reasonable, he largely brackets who controls the procedures. If we did identify fair principles (via something like the Veil of Ignorance—which I believe is not practical—or overlapping consensus), who would implement those principles? Gabriel warns against “value imposition and domination” and emphasizes “non-domination” as a criterion. I believe that his thinking would successfully cash out in the deployment of real systems if we modeled relatively benign actors who want to do the right thing.

I also worry that fair procedures, while nice in theory, may be difficult to implement when people have incommensurable fundamental commitments. Even such procedures may be difficult to resolve when people come to the table with fundamentally incompatible views and foundations—in this case, control of procedures becomes even more consequential because there is no neutral ground.

Consider economic distribution: a frequently posed question is, should AI-driven automation require redistributing gains to displaced workers? A libertarian might say: no, because property rights are invaluable. A socialist might say: yes, automation’s benefits should be socialized. On content moderation, a US conservative might say AI should not filter “hate speech” because this violates free speech absolutism. A European social democrate might say it should, because dignity trumps unlimited speech; an Islamic scholar might also say yes, but use a different definition of prohibited speech.

The worry about incommensurable values points to a deeper question about whether alignment frameworks assume too much universality or aggregate coherence in values. I think an interesting perspective on Gabriel’s thinking comes from Yuk Hui. Known for his ideas of “technodiversity” and “cosmotechnics,” Hui argues that there is not one universal technology, but multiple cosmotechnics embedded in different cosmologies and philosophies.

A way I like to frame this question is: had the technologies we know today (in the US, for me) developed in a society with totally different social and political foundations, how would they look? In what ways would we admit them into our lives? In a move against Heidegger, Hui argues for different relationships between technology and cosmos/nature depending on cultural-philosophical traditions.

For example, in this LARB interview, Hui discusses 天下 (tianxia), a cosmotechnics that legitimized government through cosmic alignment, not procedural fairness. As cosmotechnics of this type is not possible in a time with no conception of “Heaven”—that is to say, if we do not recognize the heavens as a morally legitimizing power, we cannot properly conceive of the values that would hold sway over us under 天下. This sort of recognition influences how we construct meaning from experience. Under such a cosmotechnics, the question of “alignment with whose values” might not arise in the same way. If technology is understood as harmonizing cosmic and moral order (and is developed in such a way), alignment would not be a separate normative problem to solve—it belongs to the conception of technology itself.

Looking back to Gabriel’s framework—which takes liberal-democratic categories as starting points—cosmotechnics says something about what is embedded in technologies before we consider alignment, conceptual translation difficulties across frameworks that don’t share basic categories, and scope limitations of individual philosophical approaches.

For a society like ours, Gabriel’s liberal egalitarian framework feels natural and appropriate. But understanding how technologies already embed worldviews helps us see what alignment projects can and can’t accomplish. Alignment efforts work downstream of fundamental choices about how technologies are conceived and built—by the time we ask whose values a particular system should reflect, we are often already operating with a system that embeds values and choices.[^1] Tom Mullaney’s work in The Chinese Typewriter shows how the very concept we develop of a technology constrains what it can do, and for whom. In alignment, then, we work within constraints that reflect not just technical choices but entire ways of relating technology to human life—ways that procedures alone cannot remake.

An aside / post-script: I really like Hui’s use of the Kantian antinomy in this essay. (1) Technology is anthropologically universal as the exteriorization of memory, and (2) technology is not anthropologically universal because it is conditioned by particular cosmologies. (1) and (2) are each insufficient, and their tension is productive.

[^1] to be clear, Gabriel understands and considers this in his work—he would be one of the first to say that our development of technology should be guided by inputs that represent different people’s values. even then, our imagination of technological possibilities is constrained by a history of value judgments, and we are never operating from a “neutral position.”

hypographic media?

Sat, 30 Aug 2025 00:00:00 +0000

When I started writing this post, I built a few prototypes of what something called “hypographic media” could look like. I played around with them, decided they were all paltry, and concluded that it was better to describe the thing I was trying to understand in words rather than prototypes, because I am a creature of limited imagination and because it exists in a way that I don’t believe a web demo can access.

Tom Mullaney, who has so generously let me interview him and given me some of his time, made a point over a phone call that I haven’t been able to stop thinking about. He said: “We have hypographic mediums, but no hypographic media.” Let me explain.

In the Introduction to his excellent 2024 book, The Chinese Computer: A Global History of the Information Age, Tom discusses a crisis affecting a billion Chinese people: they are forgetting how to write. This isn’t a phenomenon affecting those without access to educational resources, but rather stalks China’s elite. Wealthy, urbanized, high-income Chinese are forgetting how to write even the most basic characters. The culprit is digital writing.

Tom introduces the word hypography to name how writing is occurring in this digital era, and argues that this new “technology of the intellect” has been taking shape over the past half-century. While orthographs, the results of conventional writing (or orthography), are formed by patterns of letters that we recognize as words bearing particular meanings, hypographs are a class of writing whose job is to help search for and retrieve orthographs from memory. The sequence of letters diantmvv is a hypograph used by a Chinese “Input Method Editor” (IME), which intercepts the keystrokes and figures out what Chinese characters the user who typed diantmvv wants to produce. The letters diantmvv that our user typed do not mean anything on their own, and only serve the purpose of directing retrieval. It is only the orthograph, the Chinese characters proposed by the IME and selected by the user, that perform the labor of meaning something. In short, orthographs are writing that says something, while hypographs retrieve or produce orthographs.

Hypography needn’t be meaningless. In his conclusion, Tom states that the prompts we give to LLMs are a hypographic sequence. While Chinese IMEs merely predict a preexisting intention—a writer already knows the characters they intend to write and merely inscribe the sequence of letters used to retrieve those characters—someone prompting an LLM may not have thought of the words or passages presented to them by an LLM when they prompt it to ideate or produce an essay. In this case, it is not necessary that the hypograph has no meaning of its own (prompts mean something!), but it is true that in the production of a final artifact, the hypograph vanishes.

Working definition: hypographic media are socially realized systems where retrieval or selection traces (prompts, queries, IME codes) are preserved and read as communicative objects, rather than being discarded as inputs.

Returning to Tom’s statement, what does it mean to have hypographic mediums, but no hypographic media? Mediums are often thought of as conduits; a television or a book is a medium for the transmission of information. Generations of media scholars have discussed how mediums also actively shape what and how we perceive and organize social life. That form and environment matter more than “content” and influence our reception of the content is well-summarized in McLuhan’s dictum “the medium is the message.” Media, on the other hand, is defined by Lisa Gitelman as “socially realized structures of communication,” where the technologies and protocols come together. John Durham Peters advocates for a more expansive notion of mediums and media: everything can be a medium, and the concept of media should not be confined to channels or institutions. Rather, media are infrastructures and environments we live within.

What, then, are hypographic media? We’re putting together two concepts: hypography, a form of writing that exists at a register removed from conventional writing and, in my reading, is the immediately erased evidence of decisions made in the process of producing an artifact; media, a socio-technical system that organizes communication. Hypographic mediums, which we do have, are the technical channels or tools that transmit hypography: IME interfaces, search boxes, prompt windows. Hypographic media, by analogy, could be thought of as socio-technical systems that make hypographic traces into communicative artifacts in their own right, as opposed to disposable inputs. Rather than vanishing once the orthograph appears, the hypograph is preserved, circulated, interpreted, and institutionalized.

I think numerous forms of proto-hypographic media already exist. “Process artifacts” can be used to describe the tangible byproducts of the development of a project or system: in software development, this might take the form of design documents, meeting notes, and other reference material. This gives insight into the choices made, tradeoffs considered, and other thinking that might have factored in the production of a final artifact but do not, themselves, exist as part of that artifact.

On the other hand, Wikipedia edit histories and git commits seem related because they are parts of the production of an orthograph, but are better thought of as paratexts of orthography because they are revisions of an orthograph as opposed to inscriptions whose instrumental role is to search for or produce orthographs.

While we can spend plenty of time quibbling over boundaries, it’s more interesting to consider the stakes. Why should we care about whether hypographic media exist?

In education, we often look for evidence of thinking and understanding in legible formats: the grade, the essay, and so on. The orthograph—an essay with no accompaniment—may no longer be sufficient evidence to conclude that understanding has occurred unless produced in specific circumstances. In The New Yorker, Princeton historian D. Graham Burnett describes how he assigns students to work with ChatGPT, then reads their dialogues; he treats some of the students’ rich exchanges with the model about theology and other subjects as serious intellectual work. In reading the process (which may or may not have led to an orthographic output such as an essay), Burnett is modeling hypographic media in education: the conversation that makes an essay or other output possible is itself a communicative object used for learning and assessment.

If writing is a mode of thinking (in fact, the dominant mode of thinking in our literate culture), hypogaphy is the part that evaporates—the choices disappear while the final artifact remains. For decades, many Chinese speakers have been using hypographs to write; now, many of us seeking to write or ideate are also using hypographs to produce a final artifact that is a piece of writing or an articulation of our own thinking.

Writing, as Tom suggests, has attained a new form in hypography. Thinking, and the evidence we use to evaluate it, has changed with it. I don’t want every keystroke immortalized, and (of course) private chatbot conversations shouldn’t be public by default. The alternative is selective publicity: keep the under-text private most of the time, but make particular hypographs citable where they serve learning, review, or safety (in classrooms, science, or safety-critical systems). In these zones, hypographic media are evidence and, sometimes, the work itself.

how to lie with embeddings

Fri, 13 Jun 2025 00:00:00 +0000

Incoherence seems to me preferable to a distorting order. —Barthes

One of the ideas I thought about a lot when studying metaphysics and that continued to find me everywhere since then has been Projectivism. In our perception and understanding of the world, we attribute to it structure that it doesn’t actually have. Hume’s argument against causation used a form of this diagnosis: we often perceive two events or a series of events in sequence—B always follows A—and imagine that there is some causal relation between A and B, so that A causes B. We never observed the causal relation itself, though, so Hume thinks we have no justification for believing in that relation. All that we observe is correlation.

Visualizations are useful for understanding data, but they blur the line between what the data actually show and the patterns we project onto data. Every visualization embeds assumptions about the data’s structure, and the challenge isn’t just validating our assumptions — it’s recognizing when we’ve convinced ourselves we see meaningful patterns that exist only in our interpretation and not in reality. This distinction becomes especially dangerous when the same dataset can be made to tell completely different stories depending on our analytical choices.

Some Research Behind the Deception

Before we look at experiments, let’s understand why these visualizations can be problematic. Several landmark studies have revealed some of our interpretive biases and misunderstandings:

Visualization Mirages

Albert et al. (2018) showed that even randomly sampled crime incidents produced illusory “hot-spots,” leading participants to re-allocate police resources. Using the VALCRI (Visual Analytics for Sense-making in CRiminal Intelligence analysis) project, the authors asked participants to evaluate whether they would increase police presence in one of two city districts, along with follow-up questions about how the data influenced their decisions and whether they could justify their decisions, given tools that showed spatial and chronological distribution of crime incidents in those two districts. They were given “random condition” data randomly selected from a large set of incidents, “pattern condition” data reflecting real spatial and temporal patterns, and presented with the data in an “interactive condition” where they could interact with tools to inspect incidents from different perspectives as well as a “static condition” where they could not interact with tools.

McNutt et al. (2020) later coined the term visualization mirages for silent but significant failures that arise at any stage of the analytic pipeline.

Cleveland & McGill’s Graphical Perception Hierarchy

Cleveland & McGill’s classic experiments rank visual encodings by accuracy: position ≫ length/angle ≫ area ≫ color. Because t-SNE and UMAP re-encode high-dimensional distance as area density and color, they push viewers toward less reliable perceptual channels—making misreading almost inevitable.

Position along a common scale (most accurate)
Position along non-aligned scales
Length, direction, angle
Area
Volume, curvature
Shading, color saturation (least accurate)

Embedding visualizations face three key challenges: algorithmic sensitivity to parameter choices, method-specific trade-offs between local and global structure preservation, and the gap between what the algorithms optimize for versus what viewers need to interpret. While humans are capable of perceiving patterns and relative distances, different algorithms make different implicit choices about which aspects of high-dimensional structure to prioritize—choices that can dramatically change the story the same data appear to tell.

A Few Experiments

Let’s look at how this happens through a few experiments. I’ll limit focus here to t-SNE and UMAP, two popular embedding methods. Both are powerful methods, but need to be treated with some caution, as I’ll hope to illustrate below.

Experiment 1: Finding Patterns in Pure Randomness

This experiment demonstrates what researchers call “apophenia”—our tendency to see meaningful patterns in random data. The algorithm parameters don’t just reveal structure; they impose some structure that isn’t there.

Research on ensemble perception shows a similar pitfall: observers can summarise large point clouds quickly, but their subjective confidence often diverges from ground-truth accuracy. Two examples are the survey of ensemble coding tasks by Szafir et al. (2016) and the “Regression by Eye” experiments by Correll & Heer (2017).

Wattenberg’s Distill guide to t-SNE (henceforth Wattenberg et al. (2016)) explains that the algorithm expands dense areas and contracts sparse ones—“cluster sizes … mean nothing”—and that lowering perplexity can manufacture clusters in pure noise.
UMAP’s own documentation warns that it “does not completely preserve density” and “can also create false tears in clusters” (UMAP-learn docs).
Kobak & Berens (2019) provide biological case-studies where such artefacts mislead interpretation, and Chari & Pachter (2023) show that, in large single-cell benchmarks, neighbor-overlap often falls below 0.3.

Experiment 2: Hyperparameters matter

Watch how the same data tells different stories based on parameter choices (“Those hyper-parameters really matter,” Wattenberg et al., 2016). Wattenberg et al. (2016) showed that t-SNE is so sensitive to hyperparameters that you can make data look like it has distinct clusters or continuous structure that doesn’t exist.

In this experiment, the range is a bit small (perplexity 2-50) and some of the variance only really comes out at higher perplexity values like 100 or so. But you can still see some differences!

Experiment 3: Confirmation bias / view selection

Running embeddings until one “looks interesting” (a real thing that happens!) is a visual form of p-hacking. A 2025 CHI study on confirmation bias in dashboard “data facts” shows that analysts overwhelmingly choose views that confirm prior beliefs and ignore contradictory ones

What the Research Tells Us

The research is clear about several critical issues:

1. Cluster Assumption

Wattenberg (on t-SNE) and Kobak & Berens both demonstrate that visually separate islands can be artefacts. Users assume that visual clusters represent meaningful groups in the data, but this assumption is often violated by dimensionality reduction algorithms. As noted in the t-SNE literature, visual clusters can appear even in structured data with no clear clustering, making them potentially spurious findings.

2. Stability Illusion

Wattenberg also notes that visual “stability” can be forced by parameter tweaking without improving fidelity.

3. Narrative Fallacy

Once users see a pattern, they create stories to explain it. Two papers from Cindy Xiong (2023 and 2019) and collaborators offer complementary findings:

Belief-biased estimates (Xiong 2023) Viewers who expect a relationship between two variables over- or under-estimate r-values by ≈0.1.
Causal story-telling from correlation (Xiong 2019) With the same correlational dataset, 33–39 % of study participants who saw a two-bar summary, and ~20 % who saw a scatter plot, wrote explanations that implied causation, despite being reminded that “correlation ≠ causation.” High aggregation (two bars) and grouped encodings produced the strongest causal ratings, while fully disaggregated scatter plots produced the weakest.

Implications for Practice

Here are a few evidence-based recommendations:

For Creators of Embeddings

Always show multiple parameter settings - Single visualizations are misleading by default
Report distance preservation metrics - Quantify how well distances are preserved
Use stability analysis - Show how consistent patterns are across runs
Document all preprocessing - Feature scaling and selection dramatically impact results
Provide interaction - Let viewers explore the parameter space themselves

For Consumers of Embeddings

Ask these critical questions:

What parameters were used? Were they chosen before or after seeing the results?
How stable are these patterns across different runs?
What preprocessing was applied to the data?
How well are distances preserved from the original space?
What would the visualization look like with different parameters?

Building Better Practices

The solution isn’t to abandon embedding visualizations entirely—they can be useful exploratory tools when used responsibly. The key is to treat them as hypothesis generators, not hypothesis confirmers.

Validate!

Always validate patterns found in embeddings using other methods:

Statistical tests in the original high-dimensional space
Domain expert evaluation
Predictive modeling to test if clusters are meaningful
Stability analysis across multiple runs and parameters
User studies comparing projection methods highlight that some layouts feel trustworthy yet score poorly on objective overlap metrics—see the perception-based evaluation by Etemadpour et al., 2015.

Toward Honest Visualizations

The research makes one thing clear: our visual system and cognitive biases make us sitting ducks for embedding deceptions. We see patterns where none exist, sometimes create stories to explain randomness, and remain confident in our misinterpretations. This is important, especially if real decisions about research and resource are to be made on the basis of how people interpret these figures.

On the technical end, we need better tools that communicate uncertainty and parameter sensitivity. Interpretively, we shouldn’t think of embedding visualizations as persuasive devices but instead as exploratory tools that we treat with appropriate skepticism.

The ability to create beautiful visualizations is not the same as the ability to reveal truth. Sometimes, the most honest thing we can say about complex data is that it’s complex—and no amount of algorithmic magic will change that.

References

Albert, D. et al. 2018. Effect of Clustering Illusion during the Interaction with a Visual Analytics Environment
Chari, T. & Pachter, L. 2023. The Specious Art of Single-Cell Genomics
Cleveland, W. & McGill, R. 1984. Graphical Perception: Theory, Experimentation, and Application to the Development of Graphical Methods
Correll & Heer, 2017. Regression by Eye: Estimating Trends in Bivariate Visualizations
Kobak, D. & Berens, P. 2019. The Art of Using t-SNE for Single-Cell Transcriptomics
Li, S. et al. 2025. Confirmation Bias: The Double-Edged Sword of Data Facts in Visual Data Communication
McNutt, A., Kindlmann, G., & Correll, M. 2020. “Surfacing Visualization Mirages”
Etemadpour, R. et al. 2014. Perception-Based Evaluation of Projection Methods for Multidimensional Data Visualization
Szafir et al., 2016. Four types of ensemble coding in data visualization
UMAP-learn documentation, Using UMAP for Clustering
Wattenberg, M., Viégas, F., & Johnson, I. 2016. How to Use t-SNE Effectively
Xiong, C. et al., 2019. Illusion of Causality in Visualized Data
Xiong, C. et al., 2023. “Seeing What You Believe or Believing What You See? Belief Biases Correlation Estimation”

my favorite essays/writing about classical music

Sat, 03 May 2025 00:00:00 +0000

I haven’t found a great compendium of pieces/writing about classical music, so I thought I’d put one together so I at least have this to refer to later. I also probably haven’t looked hard enough, so if you’re reading this and I’ve totally missed something I shouldn’t have missed, send me a note. Also: this is short! I’ll add to it in time!

“The Monster Concerto” — this piece and my first listen to Kirill Gerstein’s excellent recording made me obsessed with the Busoni Piano Concerto. Busoni’s Sketch of a New Esthetic of Music is also a fascinating read, for its discussion of the arts in general, for bangers like: “The greater part of modern theatre music suffers from the mistake of seeking to repeat the scenes passing on the stage, instead of fulfilling its own proper mission of interpreting the soul-states of the persons represented.” Embedded is another performance by Marc-André Hamelin:
“Apparition in the Woods” — the Sibelius Violin Concerto is one of my favorite pieces of all time. This recording of Hilary Hahn performing it with the Filharmonica Georges Enescu caused me to freeze when I first heard it because of how indignant and impassioned it felt compared to her other performances.
“Augustin Hadelich’s Bold Violin Explorations” — again Alexis Ross in The New Yorker, sorry. Agustin Hadelich is lovely.
“The Concerto Challenge” — apologized for that once, not doing it again.
(not writing) TwoSet Talks — I love their comedy, I also particularly loved their interviews with James Ehnes (who is incredible, like, watch his thumb in this) and Agustin Hadelich.
Absolutely on Music — a series of conversations between writer Haruki Murakami and the late Seiji Ozawa. I feel like I picked up a lot about Ozawa’s understanding of his own conducting style here.
A Very Private Sense of Fear: An interview with violinist James Ehnes —
“Do You Actually Know What Classical Music Is? Does Anyone?” — I really appreciate general audience writing on classical music. I wish more of my friends (and people in general) listened to it!
“How Leonard Bernstein Changed the Canon”
“Yuja Wang and the Art of Performance” — she’s so cool
(not writing) Zlatomir Fung: “I don’t want to be known for winning the Tchaikovsky” — this interview made me a huge fan, and is really just an excellent, thoughtful conversation. I wrote a little about it here!
“Yo-Yo Ma’s Musical Mind” — I mean.

attention as kernel smoothing

Thu, 09 Jan 2025 00:00:00 +0000

Fun fact: Cosma Shalizi is sort of my grand-advisor. To be precise, he was my undergraduate research advisor George Montañez's PhD advisor (I'm now remembering the tributes to Manuel Blum after my friend Sheon's profile of him: "Manuel Blum was not my advisor, he was my advisor's advisor.", "Manuel Blum was not my advisor, he was my advisor's father").

In a great blog post, Shalizi wrote about how attention is a form of kernel smoothing. I think there are a few takeaways here. One is the relationship between the concepts themselves, which I'll elaborate a bit on here. The second is that some of the most game-changing research ideas aren't "new" in the sense we like to imagine — they're not ex nihilo. But even the same idea, the same underlying mathematics, presented with a different light in a different context with different terminology and intuition... can do so much more than it might have without this new presentation.

I'll add more to this in time, but the main idea that Shalizi presents is that what was branded "attention" was really a form of kernel smoothing. Consider a kernel function $K(u,v)$ that measures how similar $u$ is to $v$, that is non-negative and maximized when $u = v$ (what's the simplest one you can think of?).

A kernel adjusted for numerical overflow/underflow looks like this:

$$ K(u,v) = \text{exp}\left(\dfrac{\mathbf{w_1}u \cdot \mathbf{w_2}v}{\sqrt{d}}\right) $$

where the vectors $u$ and $v$ are $d$-dimensional.

The kernel function can then be used as a weight in the average

$$ \sum_{i=1}^n y_i \dfrac{K(x_i,x_o)}{\sum_{j=1}^n K(x_j,x_o)} $$

In attention, $x_o$ is the query vector, the $x_i$ are the key vectors, and $y_i$ are the value vectors. In self-attention, $y_i = \mathbf{r}x_i$ for another square matrix $\mathbf{r}$.

Anyway, I like staring at pictures and diagrams sometimes, so here's a little tool to play around with attention and kernel smoothing over an input. The input sequence is something much simpler than we actually work with: it's 50 points evenly spaced from 0 to 1 where each point has a value determined by a combo of sinusoidal functions. Kernel smoothing is intuitive — attention is taking each point's 4D positional encoding and transforming it into a query/key vector, computing similarity scores between those transformed positions, then using those similarities as weight to create a weighted average of the values (sinusoidal function vals). That's just attention!

I'm not a UI Guy, so some explanation of what's going on is in order. What you're looking at is a little plot of the values produced by kernel smoothing and attention over the same input. Numerically they're not quite the same, as you should expect! I also added a little visualization of weight profiles at the middle position of the input sequence: we're seeing how much weight/attention the middle point gives every other point. In this case things don't look terribly different. Kernel smoothing weights decay symmetrically (we have a bell curve centered at 0.5); attention weights can learn more complex patterns but here we just see a flatter bell curve.

A few notes on the relationship:

Both mechanisms fundamentally are computing weighted averages, but they have different approaches to determining the weights — kernel smoothing uses a fixed function of position difference while attention learns its similarity function.
Given some of the above, kernel smoothing is explicitly local while attention can learn arbitrary similarity patterns. So a properly trained attention mechanism could mimic kernel smoothing (and they'd look the same in the plot!), but it can learn other patterns as well.

Alpa

Fri, 20 Sep 2024 00:00:00 +0000

Alpa

Alpa is a library that automates model-parallel training for large deep learning models: it automatically generates execution plans that unify data, operator, and pipeline parallelism. Its key idea is viewing parallelisms in two hierarchical levels: inter-operator and intra-operator. It designs compilation passes to automatically derive efficient parallel execution plans at each parallelism level.

Basics and Definitions

The conventional view of ML parallelization approaches splits them into three categories:

Data parallelism: partitions the data across devices and trains the model on each partition in parallel. After each worker computes parameter updates on its data split, it needs to synchronize with other workers before performing a weight update.
Operator parallelism: partitions the computatin of an operator (e.g. a matmul) along non-batch axes and compute each part of that operator in parallel across multiple devices. This strategy requires communication to fetch input data from other devices.
Pipeline parallelism: places different groups of ops from the model graph (stages) on different workers, splits the training batch into microbatches, and pipelines the forward and backward passes across microbatches on distributed workers.

While approaches such as Megatron-LM manually combine theses three parallelisms, explorations of auto-parallelization relying on this view have limitations. Using an example from the paper: if you’re already employing operator parallelism and want to introduce data parallel replicas, you have to introduce a new set of devices and figure out the optimal operator parallelism scheme within those devices.

The key difference between the types of parallelism presented is their granularity and whether they take operators to be a basic unit.

data, operator, pipeline parallelism
intra-operator parallelism: partitions ML operators along one or more tensor axes and dispatches those partitions to distributed devies. This achieves better device utilization, but has larger communication overhead (since it needs to communicate at every split/merge of partitioned operators).
inter-operator parallelism: slices the model into disjoint stages and pipelines the execution of stages on different sets of devices. With proper slicing the communication overhead can be light, but scheduling constraints cause device idle time (see work like GPipe).

Alpa’s Strategy

Naturally, then, if you’re thinking about mapping parallelism to devices in a compute cluster, it makes sense to map intra-operator parallelism to devices with high communication bandwidth, and inter-operator parallelism to devices with less communication bandwidth. As noted, it works at two key levels:

Intra-op optimization: Minimize the cost of executing a stage of the computational graph w/r/t/ its intra-operator parallelism plan on a given device mesh (a set of devices with high inter-device bandwidth, e.g. GPUs within a server).
Inter-op optimization: Minimize inter-op parallelization latency, w/r/t/ slicing the model and device cluster into stages and device meshes and mapping into stage-mesh pairs. This requires knowing the execution cost of each stage-mesh pair reported by the intra-op optimizer.

It introduces three compiler passes, as shown below:

Alpa compiler passes and runtime architecture.

Given Jax IR and a cluster config, the inter-op pass slices the IR into stages, cuts the cluster into device meshes, then assigns stages to device meshes and invokes the intra-op pass on each stage-mesh pair to determine the execution cost of its assignment. It repeatedly queries the intra-op pass and uses DP to minimize inter-op parallel execution latency and achieve the best slicing scheme.

A very simplistic version (not as described in the paper) could look like this, assuming we have certain methods/types:

def inter_op_pass(jax_ir: JaxIR, cluster_config: ClusterConfig) -> List[StageMeshPair]:
    num_operators = len(jax_ir.operators)
    num_devices = cluster_config.num_devices
    
    dp = [[float('inf')] * (num_devices + 1) for _ in range(num_operators + 1)]
    dp[0][0] = 0
    backtrack = [[[]] * (num_devices + 1) for _ in range(num_operators + 1)]
    
    # DP
    for i in range(1, num_operators + 1):
        for j in range(1, num_devices + 1):
            for k in range(i):
                stage = jax_ir.operators[k:i]
                for mesh_shape in get_possible_mesh_shapes(j - k):
                    mesh = DeviceMesh.create(mesh_shape)
                    cost = intra_op_pass(stage, mesh, cluster_config)
                    if dp[k][j-len(stage)] + cost < dp[i][j]:
                        dp[i][j] = dp[k][j-len(stage)] + cost
                        backtrack[i][j] = backtrack[k][j-len(stage)] + [StageMeshPair(stage, mesh, cost)]
    
    # return best slicing scheme
    return backtrack[num_operators][num_devices]

Intra-op optimization

The intra-op pass solves an ILP to minimize its execution cost. It uses SPMD-style intra-op parallelism to reduce its search space, since SPMD partitions operators evenly across devices and executes the smae instructions on all devices. There are a few key ideas we need to understand how the intra-op pass works:

Device mesh: a 2-D logical view of a set of physical devices, where each device in a mesh has the same compute capability.
Sharding Spec: The sharding spec defines the layout of a tensor — each dimension is sharded or replicated across devices. The layout of an $N$-dimensional tensor is described as $X_0X_1…X_{N-1}$, where $X_i \in {S,R}$ indicates partitioning or replication for the $i$th dimension. For a 2-dimensional tensor, a spec $SR$ means it is row-partitioned. The authors also introduce a superscript to $S$ to denote device assignment: $S^0$ indicates partitioning along the 0th axis of a mesh, while $S^{01}$ indicates partitioning along both mesh axes.
Resharding: If the input tensor of an operator doesn’t satisfy the sharding spec of the parallel algorithm chosen for that operator, a layout conversion (resharding) is necessary.
Parallel algorithms of an operator: Based on an analysis of an operator’s expression (e.g. a batched matmul), we can work out possible algorithms — these involve a parallel mapping, an output (sharding) spec, input specs, and communication cost (e.g. all-reduce communication).

The ILP formulation for this pass uses an objective function that minimizes the sum of compute and communication costs. Its decision variables represent the choice of parallel algorithm for each op and the constraints ensure exactly one algorithm is chosen for each operator.

to work out communication costs, the authors compute the numbers of communicated bytes and divide by the mesh dimension bandwidth.
compute costs are set to zero — the authors argue that this is reasonable, since (1) we don’t allow replicated computation for heavy ops like matmul, and (2) computation costs are negligible for the lightweight operators where we do allow replication.

Here’s a heavily simplified version, assuming we have methods like get_parallel_algorithms and estimate_resharding_cost:

def intra_op_pass(operators: List[Operator], mesh: DeviceMesh) -> Dict[Operator, ParallelAlgorithm]:
    prob = pulp.LpProblem("Intra_Op_Optimization", pulp.LpMinimize)
    
    vars = {}
    for op in operators:
        for alg in get_parallel_algorithms(op, mesh):
            vars[(op, alg)] = pulp.LpVariable(f"{op.name}_{alg.input_spec}_{alg.output_spec}", cat='Binary')
    
    # objective
    prob += pulp.lpSum(vars[(op, alg)] * (alg.communication_cost + alg.computation_cost)
                       for op in operators
                       for alg in get_parallel_algorithms(op, mesh))
    
    # add constraints
    for op in operators:
        prob += pulp.lpSum(vars[(op, alg)] for alg in get_parallel_algorithms(op, mesh)) == 1
    
    # resharding costs
    for i in range(len(operators) - 1):
        op1, op2 = operators[i], operators[i+1]
        for alg1 in get_parallel_algorithms(op1, mesh):
            for alg2 in get_parallel_algorithms(op2, mesh):
                resharding_cost = estimate_resharding_cost(alg1.output_spec, alg2.input_spec, np.prod(op1.output_shape), mesh)
                prob += pulp.lpSum(vars[(op1, alg1)] * vars[(op2, alg2)] * resharding_cost)
    
    prob.solve()
    
    solution = {}
    for op in operators:
        for alg in get_parallel_algorithms(op, mesh):
            if pulp.value(vars[(op, alg)]) == 1:
                solution[op] = alg
                break
    
    return solution

Inter-op optimization

Let’s formalize inter-operator parallelism a bit more. We’ll think of the computational graph as a sequence of ops following the graph’s topological order, written as $o_1,…,o_{k-1}$. As described earlier, the operators are described into $S$ stages $s_1,\ldots,s_S$ where each stage contains operators $(o_{l_i},\ldots,o_{r_i})$ and each stage $s_i$ is assigned to a submesh of size $n_i \times m_i$, from a cluster mesh with shape $N \times M$.

The latency of executing stage $s_i$ on submesh of size $n_i \times m_i$ is written $t_i = t_{intra}(s_i,Mesh(n_i,m_i))$ — given $B$ input microbatches for the pipeline, the total minimum latency for the computation graph is given by

\begin{equation} T^* = \min_{s_1,\ldots,s_S; (n_1,m_1),\ldots,(n_S,m_S)} \left{\sum_{i=1}^S t_i + (B-1) \cdot \max_{1\leq j\leq S} {t_j}\right}. \tag{2} \end{equation}

where the first term is the total latency of all stages, and the second is the pipelined execution time for the other $B-1$ microbatches, bounded by the lowest stage. The figure below illustrates the pipeline latency:

The authors propose a DP algorithm to find $T^*$. [tk explain this better]

Here’s how you might write the inter-op pass, more or less based on the psudocode in Algorithm 1.

def inter_op_pass(G: ModelGraph, C: Tuple[int, int], B: int) -> float:
    N, M = C
    operators = flatten(G) # topological sort, yields (o1,...,ok)
    layers = operator_clustering(operators)
    L = len(layers)

    submesh_shapes = [(1, 2**i) for i in range(int(np.log2(M))+1)] + [(i, M) for i in range(2, N+1)] # {(1,1), (1,2), (1,4), ..., (1,M)} U {(2,M), (3,M), ..., (N,M)}
    
    t_intra = {}
    
    # t_intra for all possible stage-mesh pairs
    # recall that t_intra(s_i, Mesh(n_i, m_i)) is the minimum execution time of stage s_i on a mesh of size n_i x m_i
    for i in range(L):
        for j in range(i, L):
            stage = layers[i:j+1]
            for n, m in submesh_shapes:
                for s in range(1, L+1):
                    t_intra[(i, j, n, m, s)] = float('inf')
                for (n_l, m_l), opt in logical_mesh_shape_and_intra_op_options(n, m):
                    plan = intra_op_pass(stage, DeviceMesh((n_l, m_l)), opt)
                    t_l, mem_stage, mem_act = profile(plan)
                    for s in range(1, L+1):
                        # check whether the required memory fits the device memory
                        # in the 1 fwd 1 bwd schedule, max_stored_activations = s and this reduces to Equation 5
                        if mem_stage + max_stored_activations * mem_act <= mem_device:
                            if t_l < t_intra[(i, j, n, m, s)]:
                                t_intra[(i, j, n, m, s)] = t_l

    T_star = float('inf')
    for t_max in sorted(set(t_intra.values())):
        if B * t_max >= T_star:
            break
        
        F = {}
        F[(0, L+1, 0)] = 0
        
        for s in range(1, L+1):
            for l in range(L, 0, -1):
                for d in range(1, N*M+1):
                    F[(s, l, d)] = float('inf')
                    for k in range(l-1, -1, -1):
                        for n, m in submesh_shapes:
                            if n * m <= d and t_intra[(k, l-1, n, m, s)] <= t_max:
                                F[(s, l, d)] = min(F[(s, l, d)], 
                                                   F[(s-1, k, d-n*m)] + t_intra[(k, l-1, n, m, s)])

        T_star_t_max = min(F[(s, 0, N*M)] for s in range(1, L+1)) + (B - 1) * t_max
        if T_star_t_max < T_star:
            T_star = T_star_t_max

    return T_star

Ideas

The paper notes five different limitations.

It doesn’t handle cross-stage communication cost since that communication cost is small. This is typically true for sequential architectures, but not necessarily for models w/ skip connections or dense connectivity patterns.
Doesn’t optimize for # of microbatches B. We could adapt the DP formulation to handle this, adding DP as another dim to the optimization. This bumps the complexity of DP to $O(L^3 NMB_max)$ vs. $O(L^3NM)$.
Doesn’t consider more dynamic schedules for pipeline parallelism. I think you could also introduce schedules (e.g. GPipe) into the DP formulation — but more complexity!
Doesn’t optimize for overlapping computation and communication. Not sure how exactly this would fit into Alpa, but fine-grained scheduling is always an option: create an op dependency graph, assign each op to resources, create a time-based schedule, then schedule communication ops to overlap w/ computations that don’t depend on the communication outputs.
Only handles static computational graphs. Something with JIT compilation — idk exactly what this would look like!

AsyncTP

distributed compilation and remote execution

Fri, 13 Sep 2024 00:00:00 +0000

Motivations

Compile times, especially for large projects, can be very long. Most of us don’t worry about this too much, but it’s a big deal for things like high energy physics experiments, where, according to this paper by Rosen Matev, compiling from scratch can easily take 5+ hours even when using an 8-core VM. Incremental builds don’t help much if you change something like a widely-used header file.

As Matev explains, high energy physics experiments typically have large software codebases written in C++. The LHCb physics software stack is fairly complex, and is depicted below:

LHCb physics software stack

Here’s how Matev explains the difficulties:

Developing on such a codebase can prove difficult due to the amount of resources required for building. Typically, development happens on hosts that are readily available but often not very powerful, such as LXPLUS, 4-core CERN OpenStack virtual machines (VMs), desktop PCs. After building once, incremental builds are supported, but often the increments are very large due to modifications in headers or when switching branches. To combat the limited resources, custom tools exist for working on partial checkouts on top of released project versions [2], however they are not well suited for other than trivial changes due to, for instance, the lack of good support for standard git workflows, or the need to checkout a large chunk of the stack.

distcc

In this paper Matev describes the distributed compilation server distc. The schematic below illustrates how it works at a high level — distcc allows offloading compilation to voluntary or dedicated resources. This isn’t really a compiler — it’s a system that lets a client offload compilation to available resources. I’ll describe the design below, but it’s work noting that this is living in CMake world. If you’ve heard of Bazel, its remote execution feature can achieve similar results.

distcc Schematic

It works in two modes:

Standard Mode: A client runs the C++ preprocessor locally, and preprocessed output is sent to the server for compilation. This mode is limited by local preprocessing speed.
Pump Mode: This uses an “include server” on the client, which analyzes source files to find header dependencies. Source files and necessary headers are sent to the server, then preprocessing and compilation happen on the server. This mode asumes system headers are identical on the server and client.

There are a few other important components:

Servers use GSSAPI authentication and clients authenticate using Kerberos tokens.
distcc uses the Ninja build system with 100 parallel jobs; “pools” limit concurrenty for non-distirbutable tasks like linking.
Uses the CernVM File System to share LCG releases (CERN lingo: a set of ~200 software packages consistently built together).

distcc’s Distributed Cache

Sharing caches is hard — developers might use different directories and info contained in lookup keys might contain absolute paths, e.g. when including debug symbols. sccache natively supports remote storage such as S3 or Redis — like ccache, it’s used as a compiled wrapper but has a client-server architecture for efficient communication with the remote storage. Its lookup key determination is also more simple than ccache — this means fewer cache hits, offloading the cache but not sharing it between developers.

Bazel and remote execution

Bazel’s remote execution achieves something like that DistCC was built to do: you can distribute build and test actions across multiple machines, e.g. in a datacenter. Bazel goes beyond DistCC in its ability to scale entire build processes (including non-compilation tasks).

Misc

There was what looks like a Stanford class project from 2014 called “sdcc: Simplf Distributed Compilation” — its goal was to be language-agnostic, unlike distcc, and work without meta information.

Google has a patent on Distributed JIT compilation. I haven’t read the whole thing yet, but at a high level it sounds like what you’d expect: a client platform sends a first request message to a dedicated compilation server, compiling the bytecode sent in that first message into something the client can execute, notifying the client, then sending a second message from the client to the server requesting the instructions and accessing the repository to move the instructions from the server to the client platform.

Icecream describes itself as a distributed compiler with a central scheduler to share build load. Again, not an actual compiler: it takes compile jobs from a build and distributes them to remote machines. It’s based on distcc but uses a central server to dynamically schedule compile jobs to the fastest free server.

Google’s distributed compiler service for open-source projects like Chromium and Androis. It’s basically distcc+ccache.

Jussi Pakkanen has a few notes on architecture for a distributed compilation cluster here.

Simplifying Dependent Reductions in the Polyhedral Model

Mon, 02 Sep 2024 00:00:00 +0000

I read this really neat paper recently called “Simplifying Dependent Reductions in the Polyhedral Model” by Cambridge Yang, Eric Atkinson, and Michael Carbin. I’ll give a rundown of the main ideas: what (dependent) reductions are, why they matter (especially in ML), how this paper uses the polyhedral model to simplify them, and why that matters.

(Dependent) Reductions and why they matter:

A reduction is a pretty familiar concept: reductions combine a set of values into a single result. A reduce-add takes a list of numbers and returns their sum. A reduce-multiply takes a list of numbers and returns their product. This shows up everywhere in numerical computing, and often in ML.

Optimizing these reductions, then, can dramatically speed up workloads.

Ordinary reductions, like summing an array or finding its maximum value, have some appealing properties that make them relatively easy to optimize and parallelize: these include associativity, commutativity, and independence. Compilers and runtime systems can exploit these properties to perform parallelization, vectorization, loop unrolling, and tree-based reductions.

Dependent reductions are a different beast entirely: each step of a dependent reduction depends on the result of the previous step. This shows up in many places in numerical computing and ML, especially in convolutions, matrix multiplications, and some recurrent neural networks. The paper we’ll discuss here uses a an example of the prefix sum, but I’ll mention some examples more relevant to ML in this intro:

SGD with Momentum

The momentum update can be expressed as:

\[\begin{aligned} v_t &= \beta \cdot v_{t-1} + (1 - \beta) \cdot g_t \\ \theta_t &= \theta_{t-1} - \alpha \cdot v_t \end{aligned}\]

Where:

$v_t$ is the velocity at time $t$
$\beta$ is the momentum coefficient
$g_t$ is the gradient at time $t$
$\theta_t$ is the parameter at time $t$
$\alpha$ is the learning rate

This can be expressed as a dependent reduction:

v = 0
θ = initial_params

for t in range(num_iterations):
    g_t = compute_gradient(θ)
    v = beta * v + (1 - beta) * g_t  # Dependent reduction
    θ = θ - alpha * v  # Dependent reduction

where each update to $v$ and $\theta$ depends on its previous value, forming a chain of dependent reductions.

Optimizing a reduction like this could produce significant performance improvements in training models.

Self-Attention

Attention mechanisms aren’t really implemented as dependent reductions, but certain variants (e.g. inremental self-attention in autoregressive decoding) look conceptually similar to dependent reductions.

When we compute attention scores incrementally, as below, we can express them as dependent reductions:

def incremental_self_attention(queries, keys, values, t):
    attention_scores = np.zeros((t+1, t+1))
    context_vector = np.zeros(d_model)
    
    for i in range(t+1):
        for j in range(i+1):
            attention_scores[i, j] = dot_product(queries[i], keys[j])
        
        attention_weights = softmax(attention_scores[i, :i+1])
        
        for j in range(i+1):
            context_vector += attention_weights[j] * values[j]  # Dependent reduction
    
    return context_vector

Here, the context_vector is computed incrementally, with each step depending on the previous steps.

The Polyhedral Model

The polyhedral model is a powerful technique in compiler optimization — its insight is that since programs spend most of their time in loops, it’s useful to develop a simple way to express information about loops that allows us to reason about program behavior and perform optimizations. The representation offered int he polyhedral model lets compilers analyze loopnests and dependencies, automatically find opportunities for optimization, and apply transformations to improve performance (e.g. parallelization or improving memory access patterns).

I hope to write a fuller intro in another post, but for the purposes of explaining this paper I’ll give a brief intro to how this works and show how it can be used for optimizations like reordering loops. You can understand the main heuristic algorithm in this paper without too much detail on the polyhedral model, but it’s helpful to understand it as motivation.

Some Intuition

I’ll describe the model in a bit more technical detail below, but the following might help build some intuition: the polyhedral model allows us to think about loopnests as polyhedra (hence the name) defined by the bounds of the loops that make up the loopnests (technically in more than three dimensions these are called polytopeds and the polyhedral model is also referred to as the polytope method).

If we consider a doubly-nested loop:

for (i = 0; i < N; i++) {
  for (j = 0; j <= i; j++) {
    // do something
  }
}

we can think of each iteration of the loop as a point $(i,j)$ in 2D space. The bounds of the loops define a shape containing all these points — a polygon in 2D, a polyhedron in 3D, and polytopes in higher dimensions. The loop bounds define a set of points that are valid iterations of the loop, which in this case looks like a triangle:

Simple loop bounds example

With a figure like this, dependencies can be represented as arrows between points. I’ll use a similar example again below with some more detail and introduce a formalization of the polyhedral model.

Representing Loops

A loop can be represented as a set of constraints on the possible values of the loop variables. For example, the loop:

for (i = 0; i < N; i++) {
    for (j = 0; j <= i; j++) {
        A[i] += B[j];
    }
}

can be represented as the constraints:

Loop iteration space example

Each point in this polyhedron represents one execution of the innermost statement (A[i] += B[j]). The coordinates of the point correspond to the values of i and j for that execution.

Formally, we define a polyhedral set as:

$P = [p] \rightarrow \{[x] : M \cdot [x, p, 1]^T \geq 0\}$

Where:

$[p]$ is a vector of parameters (like $N$ in our example)
$[x]$ is a vector of variables (like $i$ and $j$)
$M$ is a matrix defining the inequalities

In our example:

\[P = [N] \rightarrow \{[i, j] : \begin{bmatrix} 1 & 0 & 0 & -1 \\ 0 & 1 & 0 & 0 \\ -1 & 1 & 0 & 0 \\ 1 & 0 & -1 & 0 \end{bmatrix} \cdot \begin{bmatrix} i \\ j \\ N \\ 1 \end{bmatrix} \geq \begin{bmatrix} 0 \\ 0 \\ 0 \\ 0 \end{bmatrix} \}\]

To return to the simple example we’ve been using, we can describe the triangle with inequalities:

$0 \leq i < N$
$0 \leq j \leq i$

and in the polyhedral model, we can represent this as:

\[P = [N] → {[i, j] : 0 \leq i < N \text{ and } 0 \leq j \leq i}\]

This polyhedral set defines all valid points $(i,j)$ for our loop — this is also called the iteration space of the loop. In our loop, each A[i] += B[j] operation depends on the previous iterations with the same $i$. We can represent this as arrows in our diagram.

Polyhedral reduction example with dependencies

To compute A[2], for instance, we need to perform all the additions for i=2 in order from j=0 to j=2.

In the polyhedral model we also care about schedules, which assign an execution time to each point in the iteration space. Our original loop implements the simple schedule:

\[\Theta(i, j) = [i, j]\]

meaning to execute in the order of $i$, then $j$.

This schedule can be represented more formally as a scheduling matrix. A scheduling matrix $\Theta_S$ for a statement S maps each point in the iteration space to a vector of time coordinates. For our schedule above, the scheduling matrix would look like this:

\[\Theta_S = \begin{bmatrix} 1 & 0 & 0 \\ 0 & 1 & 0 \end{bmatrix}\]

This matrix, when applied to a point $(i, j, 1)$ (we add 1 as a constant term), gives us the execution time:

\[\begin{bmatrix} 1 & 0 & 0 \\ 0 & 1 & 0 \end{bmatrix} \begin{bmatrix} i \\ j \\ 1 \end{bmatrix} = \begin{bmatrix} i \\ j \end{bmatrix}\]

Each row of the scheduling matrix corresponds to one dimension of our execution time.

If we had a different schedule:

\[\Theta_S = \begin{bmatrix} 0 & 1 & 0 \\ 1 & 0 & 0 \end{bmatrix}\]

This would swap our loop order, executing in order of $j$, then $i$. Here, the first row $\Theta_{S,1} = [0, 1, 0]$ determines the first time coordinate ($j$), and the second row $\Theta_{S,2} = [1, 0, 0]$ determines the second time coordinate ($i$).

One observation we might make is that if we reordered the original loop’s points (where we executed in the order of $i$, then $j$) in a way that still respects the dependencies, we might get a faster program (perhaps by improving memory locality, for instance).

In this case, we might realize that all the B[j] values could be summed up once, instead of repeating the sum for each i. So we write the optimized version:

sum = 0;
for (j = 0; j < N; j++) {
    sum += B[j];
    A[j] = sum;
}

In the polyhedral model, this optimization is achieved by finding a new schedule that reorders the computation while respecting dependencies.

From Independent to Dependent Reduction

As an example of a dependent reduction, the paper introduces the prefix sum — it’s exactly the loopnest we saw before:

for(i = 0; i < N; i++)
    for(j = 0; j<=i; j++)
        B[i] += A[j]

representing the summation

\[B[i] = \sum_{j=0}^{i} A[j] \quad \forall i, \ 0 \leq i < N \tag{1}\]

We would ordinarily optimize this with something called the Simplification Transformation (ST), which codifies profitably reusable computations into a set of reuse vectors. Intuitively, these vectors point from one iteration to another iteration that performs some of the same computation. Knowing about the shared computation between different loop iterations, we can restructure code to perform that shared computation once and reuse the results instead of recomputing.

In (1), the reuse vector $[1,0]^{T}$ denotes the shared computation changing $i$ to $i+1$ and $j$ to $j+0$. ST, given an equational statement like (1) and a reuse vector, transforms the statement into a set of statements that are semantically equivalent to the original statement, but reuse shared computation. Given (1) and $[1,0]^{T}$, ST would transform (1) into the following set of statements:

\[B[0] = A[0] \tag{2a}\] \[B[i] = B[i-1] + A[i] \quad \forall i, \ 1 \leq i < N \tag{2b}\]

which has complexity $O(N)$, better than the naive $O(N^2)$ complexity of the original loop, which it achieves by setting a base case and reusing computation (like a simple DP).

We could have instead used the reuse vector $[-1,0]^{T}$ to denote the shared computation changing $i$ to $i-1$ and $j$ to $j-0$, resulting in the following set of statements:

\[B[N-1] = \sum_{j=0}^{j<N} A[j] \tag{3a}\] \[B[i] = B[i+1] - A[i] \quad \forall i, \ 0 \leq i < N-1 \tag{3b}\]

which also has complexity $O(N)$.

To contrast with the prefix sum, which can be optimized rather easily since the input array A is not modified during the computation of B, the paper considers the following dependent reduction:

\[B[i] = \sum_{j=0}^{j \leq i} A[j] \quad \forall i, \ 0 \leq i < N \tag{4a}\] \[A[i+1] = f(B[i]) \quad \forall i, \ 0 \leq i < N-1 \tag{4b}\]

Now the reduction in (4a) is dependent because the value of the reduction $B[i]$ depends on the set of values ${A[j] \mid j \leq i}$ while $A[i]$ depends on the previous value of the reduction, $B[i-1]$. While ST works for simpler examples, applying it to dependent reductions introduces new dependencies — together with the program’s existing dependencies, the resulting program might end up with a dependence cycle.

For instance, if we had applied the reuse vector $[-1,0]^{T}$ to (4), we would have obtained a program with statements (3a), (3b), and (4b). This forms the dependence cycle $B[N-1] \rightarrow A[N-1] \rightarrow B[N-2] \rightarrow B[N-1]$. The reuse vector $[1,0]^{T}$ would have produced a valid program consisting of equations (2a), (2b), and (4b).

To illustrate valid and invalid reuse directions, let’s look at Figure 1 from the paper, which shows the iteration space of the prefix sum loop:

Iteration space of the prefix sum loop

The polyhedron with round dots at the top represents the iteration domain of the reduction statement B[i] += A[j] (each round dot denotes an iteration instance of the statement). The elements A[0] through A[4] on the right are the array elements A[j] to be accumulated into B[i]. The bottom polyhedron with squares it eh iteration domain for the statement A[i+1] = f(B[i]), while the middle polyhedron with diamonds is an additional polyhedron that the author’s technique inserts into a program’s polyhedral representation to denote the completion of each reduction B[i].

Each arrow represents a data dependency in iteration space — an arrow from $a$ to $b$ means that $a$ needs to execute before $b$. Figures 2 and 3, below, show the correct and incorrect optimizations of our dependent prefix sum with the reuse vectors $[1,0]^{T}$ and $[-1,0]^{T}$, respectively.

Correct optimization of the dependent prefix sum with the reuse vector $[1,0]^{T}$

Incorrect optimization of the dependent prefix sum with the reuse vector $[-1,0]^{T}$

An Integer Bilinear Program and a Heuristic

So, we’ve seen how the polyhedral model represents loops and how the Simplifiation Transformation (ST) can correctly and incorreclty optimize reductions. Extending techniques like ST to dependent reductions introduces the problem we saw above: how do we choose reuse vectors to simplify programs so that without introducing dependency cycles?

In the paper, the authors formulate their optimization as an Integer Bilinear Program, considering all constraints and looking for the best solution. To understand how they formulate the optimization problem, consider again the structure of polyhedral representations.

Remember our triangular iteration space for the prefix sum — that entire triangle is what we call the “domain” of our reduction. We can consider its components:

The entire triangle (2D face), our full iteration space.
The three edges of the triangle (1D faces), where one loop variable reaches its minimum or maximum value.
The three corners of the triangle (0D faces), where both loop variables reach their minimum/maximum values.

Each of these is called a face of the polyhedron. In the polyhedral model, a face is created when we turn some of our inequalities into equalities. For example, the bottom edge of our triangle is the face where j = 0. The paper defines a face of a polyhedral set as that polyhedral set with a subset of (potentially empty or all) inequalities changed to equality constraints.

The key insight uses this notion: instead of trying to optimize the entire reduction at once, we can think about optimizing each face separately. For each face, we want to find a “reuse vector” - a way to reuse computation from one point to another within that face.

Revisiting the prefix sum example and its iteration space:

Loop iteration space example

we can consider faces and their corresponding reuse vectors:

The entire triangle: Here we applied the reuse vectors $[1,0]^{T}$ and $[-1,0]^{T}$.
The bottom edge: This represents the base case of our sum.
The diagonal edge ($j = i$). This represents the last addition for each $B[i]$.

With this in mind, we can now formulate our optimization problem:

Variables: We have a reuse vector for each face of each reduction’s domain. These are integer variables because they represent steps in our iteration space.
Constraints: a) Our chosen reuse vectors must be valid: they should preserve dependencies and actually reduce complexity. b) The resulting program must have a valid schedule: we can’t introduce dependency cycles.
Objective: We want to minimize the overall complexity of our transformed program.

The paper formulates the scheduling problem based on previous work, particularly Pouchet et al. [2011], which ensures that the schedule respects all dependencies in the program. I won’t spend time here since, as the authors state

we only use their scheduling formalization to support our formalization of the dependent reduction scheduling problem in Section 5. Our heuristic algorithm in Section 6.2 does not require a schedule to have been computed using their scheduling formalization. Our algorithm instead relies on a sequential schedule that can be computed via any means, including a scheduling algorithm that adopts realistic scheduling constraints.

With the per-face application view of ST, the authors give a formulation of Simplifying Dependent Reductions (SDR) — their optimization problem minimizes the complexity of $prog’$, a version of $prog$ transformed by a composition of STs applied to each face. The reuse vectors for each ST must lie in a set that presents sharing ($S$), satisfies an inverse condition, and reduces complexity. Finally, there must exist a schedule $\Theta$ that satisfies the dependencies in $prog’$.

\[\begin{aligned} \text{minimize} \quad & \text{complexity}(\text{prog}') \\ \text{subject to} \quad & prog' = (ST_{f_1,\vec{r}_1} \circ \ldots \circ ST_{f_n,\vec{r}_n})(prog) \\ & \vec{r}_i \in S(f_i.\text{stmt}) \cap I(f_i.\text{stmt}) \cap L(f_i) \quad \forall i \in \{1...n\} \\ & \exists \text{ schedule } \Theta \text{ of } prog', \text{ s.t. } \Theta \text{ satisfies dependence}(prog') \\ \text{given} \quad & prog, \text{dependence}(prog) \\ \text{variables} \quad & \vec{r}_1,...,\vec{r}_n \end{aligned}\]

Where $ST_{f_i,\vec{r}_i}$ is the Simplification Transformation applied to face $f_i$ with reuse vector $\vec{r}_i$, and $S$, $I$, and $L$ are sets of constraints on the reuse vectors.

The Heuristic Algorithm

Solving the full integer bilinear program is computationally expensive, so the authors propose a heuristic algorithm:

Start with an initial sequential schedule for the program. For each face of each reduction’s domain:

Choose any valid reuse vector.
If this reuse vector is inconsistent with the initial schedule, negate it (if the negation is also valid).
Apply the Simplification Transformation with this reuse vector.

The key insight is that any valid reuse vector (or its negation) will be consistent with some sequential schedule. This allows for a much simpler optimization process.

We might write the algorithm like this:

def optimize_dependent_reductions(program):
    initial_schedule = compute_sequential_schedule(program)
    for reduction in program.reductions:
        for face in reduction.domain.faces:
            reuse_vector = choose_valid_reuse_vector(face)
            if not consistent_with_schedule(reuse_vector, initial_schedule):
                reuse_vector = -reuse_vector
                if not is_valid_reuse_vector(reuse_vector, face):
                    continue  # Skip this face if negated vector is invalid
            apply_simplification_transformation(face, reuse_vector)
    return program

This heuristic approach is proven to be optimal for reductions with inverse operators (like addition/subtraction), which covers many practical cases. This post is already long, so I’ll save more details on that for a (potential) future post.

Code!

I wrote up a little implementation of this heuristic algorithm that works for the basic prefix sum example and a dependent loop reduction — you’ll find it here!