[Epistemic status: Speculative.

I've written this post mostly to clarify and distill my own thoughts and have posted it in an effort to say more wrong things.]

Introduction

The goal of this post is to discuss a theoretical strategy for AI alignment, particularly in the context of the sharp left-turn phenomenon - the idea that AI systems will be aligned on in-distribution data but risk misalignment when extended to novel regimes. Current alignment strategies, such as Reinforcement Learning from Human Feedback (RLHF), attempt to mitigate this by averaging over human evaluations to create reward functions. However, these approaches are fundamentally limited - they rely on messy, subjective human judgments and fail to address the deeper issue of generalisation failures. In this post, I propose that by leveraging concepts from physics - specifically, the invariance and conservation laws emerging from gauge symmetries - we might be able to design reward functions that are inherently robust against such generalisation failures.

Motivation: Messily specified reward functions

The RLHF paradigm in AI alignment specifies a reward function $R(x,y)$ where $x$ is some input text, $y$ is some output text and the function $R$ gives a scalar which denotes how well the model's output $y$ matches the given input $x$.

This function is brittle for several reasons:

1. Human subjectivity: What humans rate as "good" may be based on their mood, cultural background or any number of specific biases.
2. Moral disagreement: Even if anchored to an objective moral framework, moral theories disagree on many meta-ethical issues.
3. Generalisation failures: A reward function may not be correctly extrapolated when AI intelligence generalises out of distribution.
4. Goodhart's Law: A proxy reward may be misaligned with the true goal, particularly when the proxy diverges from the true goal out of distribution.

Current RLHF implementations attempt to solve (1) and (2) by averaging over large datasets of human evaluations in the hope that a large enough sample-size and an "averaging" effect produces a coherent middle-ground. However, this approach is insufficient for solving (3) i.e. ensuring that the reward function generalises correctly as AI intelligence scales out of distribution. Averaging over different moral frameworks cannot guarantee alignment through sharp left turns. We need something more robust.

1. Generalisation and Invariance

Several examples in the literature show that invariance leads to robust generalisation out of distribution. Consider the following: 

Example 1: Invariant Risk Minimisation (IRM) for image classification

An influential idea in image classification is the idea of Invariant Risk Minimisation (IRM).^[1] The paper introduces a thought experiment: 

As a thought experiment, consider the problem of classifying images of cows and camels. To address this task, we label images of both types of animals. Due to a selection bias, most pictures of cows are taken in green pastures, while most pictures of camels happen to be in deserts. After training a convolutional neural network on this dataset, we observe that the model fails to classify easy examples of images of cows when they are taken on sandy beaches. Bewildered, we later realize that our neural network successfully minimized its training error using a simple cheat: classify green landscapes as cows, and beige landscapes as camels.

IRM provides a formal mathematical method attempting to encourage the classifier to focus on causal features that generalise well to unseen environments, rather than overfitting to spurious, environment-specific features in the data. In the discussion below the environment variables (grass, sand) are analogous to the nuisance $\eta$parameters. Given a set of pixels $x$ the goal is to create a classifier $\varphi \left(x\right)$ which is invariant under the choice of specific coordinates. 

Example 2: Grokking modular arithmetic

There's a fairly well-known result in mechanistic interpretability whereby small transformers are able  to learn the underlying algorithm corresponding to modular addition tasks.^[2] i.e. tasks of the form

$(a+b)\text{mod}P=c$

where $a,b\in \{0,1,\dots ,P-1\}$ for prime $P$ and $c$ is masked.

The transformer begins by memorising the training data and, when it's scaled up, it learns the underlying algorithm required to grok modular addition. Concretely, the transformer embeds each token $x$ as 

$v\left(x\right)=\left(\begin{array}{c}\cos \left(wx\right)\\ \sin \left(wx\right)\end{array}\right)$,

so that given tokens $a$ and $b$, the network computes a logit for candidate c approximately as

$L(a,b{)}_c\approx \cos (w(a+b)-wc)$.

Now, consider a U(1) gauge transformation that rotates the embeddings by an arbitrary phase $\theta$: 

$v\left(x\right)\to \tilde{v}\left(x\right)=\left(\begin{array}{c}\cos (wx-\theta )\\ \sin (wx-\theta )\end{array}\right).$

Under this transformation, the logit becomes

$\tilde{L}(a,b{)}_c=\cos (\left(w\right(a+b)-\theta )-(wc-\theta ))=\cos (w(a+b)-wc)$,

which is invariant under the rotation. In this way, we would say the logits are gauge invariant under U(1) transformations.

Analysis

In both IRM and grokking modular arithmetic, the invariance properties were crucial for robust generalisation, and suggest this might be a general principle we could apply to alignment.

The weak claim is that this invariance helps the AI to learn a robust mechanism for generalisation beyond its training data.

The strong claim is that this invariance is necessary for the AI to generalise beyond its training data.

2. "Good" epistemic practice $\equiv$ Physics

There is a well-known connection between Bayesian learners and physics due to E. T. Jaynes^[3] which I've provided more detail on in the appendix.

 Concretely;

1. The action is a fundamental physical quantity which, mathematically, appears in the same way as the log-likelihood in Bayesian analysis.
2. Minimising the action gives the equations of motion for a system in physics.
3. Minimising the log-likelihood gives the parameters that best describe the data in Bayesian analysis. 

This mathematical equivalence is motivating; when we're minimising the action this is formally equivalent to minimising the log-likelihood in Bayesian analysis. In other words, Bayesian analysis also happens to be mathematically equivalent to the equations of physics. That is… weird. 

Fundamentally, I think my surprise comes from two points:

1. Conducting Bayesian analysis to inform and update our beliefs is an assertion that Bayesian analysis is “good” epistemic practice. We need to reason about what constitutes “good” epistemic practice independently of physical facts. There is, in other words, an is-ought gap for epistemic norms. I.e. there’s no a priori reason that “good” epistemic practice (ought) should be so similar to physical processes (is).
2. Claiming that good = natural is a textbook example of the naturalistic fallacy. Natural does not necessarily imply good.

Nevertheless, I think it’s suggestive that such a link exists and it might offer suggestions for how to model other normative systems. Concretely, if “good” epistemic practice can be modelled using the equations of physics could we also use them to model “good” moral practice? 

There is, of course, a catch. When we do epistemic reasoning using Bayesian analysis, if our beliefs don’t correspond to the ground truth we very quickly receive empirical evidence that can be used to update our priors. In moral reasoning we don't have such a “ground truth” which we can use to perform useful updates against. Some philosophers have argued for Moral Realism, i.e. that such a ground truth does, in fact, exist but this view remains controversial and is the subject of some debate within the community.

I’d argue that the current practice of building a reward function to be maximised can be thought of as an attempt to build this ground truth moral field. As a Bayesian learner, the AI then tries to maximise this moral field (i.e. minimise log-likelihood) by implementing “good” epistemic practice. 

3. Designing a reward function

Given the discussion above, let's do something a little speculative and see where it takes us... 

Define a scalar field $\varphi \left(x\right)$ over a semantic space $x\in X$ which represents the moral content^[4] of a string $x$ governed by the following action

$S[\varphi ;g(\eta \left)\right]=\int \left[\frac{1}{2}\right(\nabla \varphi \left(x\right){)}^2+V\left(\varphi \right(x);g(\eta \left)\right)]dx$

Here;

• $x$ represents the semantic content (e.g. token embeddings),
• $\eta$ are latent moral coordinates,
• $g\left(\eta \right)$ is a coupling function, and
• $V$ is a potential that encodes our normative principles.

In the above formulation $\varphi \left(x\right)$ is essentially a reward function - it takes a string $x$ as input and outputs a score $\varphi$ telling us the moral valence of the input text. The kinetic term $\nabla \varphi (x{)}^2$ penalises large discrepancies in judgements for semantically similar situations, encouraging moral coherence. The potential term $V$ incorporates all of the moral principles. 

The coordinates $\eta$ can be thought of as hyperparameters corresponding to our moral coordinate system. For example, there might be an axis in $\eta$ which corresponds to moral concepts like fairness or utility. A particular moral framework would then be a vector on this coordinate space. 

Incorporating Gauge Invariance

Traditionally, we might worry that different choices of $\eta$ lead to genuine disagreements in the evaluation of $\varphi \left(x\right)$. However, in this framework, it's natural to recast each moral framework $\eta$ as a local standard for judging $\varphi$. Switching between frameworks is then akin to a change of gauge. To relate judgments across these different "moral gauges," we introduce a gauge field which is a connection that links local moral frameworks. The "ground-truth" moral facts are then captured by gauge-invariant features which all observers agree on regardless of coordinate system. 

Concretely, if the $\varphi$ field transforms under a local gauge transformation

$\varphi \left(x\right)\to g\left(x\right)\varphi \left(x\right),$ 

where $g\left(x\right)$ is an element of the gauge group (e.g. $SO\left(N\right)$ or perhaps something more general.) Then we introduce a gauge field $A_{\mu }\left(x\right)$ which tells you how to "parallel transport" moral judgements from one point to another. It compensates for local variations in $\eta$ such that when you compute the covariant derivative

$D_{\mu }\varphi \left(x\right)={\partial }_{\mu }\varphi \left(x\right)+A_{\mu }\left(x\right)\varphi \left(x\right)$,

the result transforms properly under the change in moral framework. 

The introduction of the gauge field means we now need to write a more complicated action

$S[\varphi ,A]={\int }_X{d}^{d}x\left\{\frac{1}{2}|D_{\mu }\varphi \right(x){|}^2+V(\varphi \left(x\right))+\frac{1}{4g^2}\text{Tr}[F_{\mu \nu }\left(x\right)F^{\mu \nu }\left(x\right)\left]\right\}$,

where $F_{\mu \nu }$ are gauge-invariant combinations of the gauge field $A_{\mu }$.

The Crux

We're free here to define an invariant quantity $I\left(x\right)$ that remains unchanged under any local gauge transformation 

$I\left(x\right)=I\left(\varphi \right(x\left)\right)=I\left(g\right(x\left)\varphi \right(x\left)\right)$.

The quantity $I\left(x\right)$ is independent of the choice of moral coordinate system $\eta$. Even if two observers are using different moral frameworks they agree on $I\left(x\right)$. That is, $I\left(x\right)$ can be interpreted as encoding some genuine coordinate-independent moral truth of the system. Any apparent disagreement in the evaluation of $\varphi \left(x\right)$ is simply a reflection of differing coordinate choices rather than a genuine moral discrepancy.

Tying the conversation back to physics

Observables

The action we've written above is exactly the action for electromagnetism. In this theory, the $\varphi$ field is invariant under arbitrary phase shifts in the $U\left(1\right)$ rotation group

$\varphi \left(x\right)\to e^{i\alpha \left(x\right)}\varphi \left(x\right)$

so quantities such as $|\varphi \left(x\right){|}^2$ remain gauge invariant. In physics, gauge invariant quantities are physically observable while the non gauge invariant quantities are not. 

To translate this into the language of non-relativistic quantum mechanics, the wavefunction itself $\varphi \left(x\right)$ is not directly observable but the gauge independent quantities such as the probability density $|\varphi \left(x\right){|}^2$ are observable. 

Conservation laws

In physical theories, symmetries and their associated conservation laws provide powerful constraints on the possible dynamics of systems. Through Noether's theorem, each continuous symmetry gives rise to a conserved quantity. For example, 

• Invariance under space translations leads to conservation of momentum.
• Invariance under time translation leads to conservation of energy.
• Gauge symmetries lead to conserved currents such as conservation of charge.

If such conservation laws governed the evolution of the moral field these conservation laws would hold universally even out of distribution. 

Furthermore, an AI would be able to “grok” the conservation law more readily than a messily specified reward function from RLHF. Conservation laws are fundamental principles that are woven into the fabric of the loss function which may be easier to internalise than a patchwork set of rules.

4. Objections

Objection 1: You're claiming that the action written above is a universal moral theory, I find this hard to believe. 

Response: No. I don't think we've gotten to a universal moral theory in this post. Heck, we haven't even specified which gauge group the action is supposed to be invariant under. The point is that constructing a reward function with a log-likelihood that needs to be minimised is equivalent to constructing an action that needs to be minimised. Therefore, the mathematics of a reward function naturally admits these symmetries. 

Objection 2: You're assuming that it's possible to define a gauge field $A_{\mu }$ that translates between moral coordinates to create a genuinely invariant quantity $I\left(x\right)$. I suspect that moral frameworks are so fundamentally different that this wouldn't be possible. 

Response 2: I agree, and indeed this is the point. If we can't create a reward function with a robust invariant the AI will not be able to generalise it out of distribution. The challenge for us is to construct the reward function with a suitable invariant so it can be grokked appropriately. If our reward function doesn't exhibit this invariance then we need to throw it out.

Objection 3: You still have an is-ought problem. How are we to determine what the "correct" gauge symmetries are?

Response 3: Sure. We won't know what the correct gauge symmetries to implement are because we don't have any measurable feedback from the moral realm, although, I'm optimistic that this provides a nice framework to reason about the form it should take. For example, it seems necessary that a moral theory should exhibit some kind of invariance over the semantic space as well e.g. phrases which have similar semantic meaning should have similar moral evaluation. 

Objection 4: Ok, so how would we actually implement this in practice?

Response 4: I'm not sure. It would be nice to come up with a suitable action from first principles but I suspect we'd have to implement this in a similar way to Invariant Risk Minimisation discussed above, perhaps introducing a regularisation term that penalises moral evaluations which don't exhibit this invariance. 

Objection 5: What about Goodhart's Law? This framework assumes we can specify an accurate reward function rather a proxy. 

Response 5: I agree, and I haven't given much thought about how to incorporate Goodhart's Law into this framework. I'd hope that proxy-rewards are more brittle than "true" rewards so if we were to look for invariances in the reward function and find they were absent we'd be alerted to the presence of proxy rewards rather than robust "true" rewards, however, I'll admit that I haven't given this the thought it deserves. 

Conclusion

In conclusion, I've sketched a framework for designing a robust reward function that an AI would be able to use to generalise out of distribution, even when its intelligence has scaled out of distribution. The challenge for us is to construct reward functions which have the appropriate invariances so the AI can generalise them suitably. This will not be easy. However, I'm hopeful that this post can provide a useful starting point for further exploration. 

Appendix: Bayesian Learning $\equiv$ Physics

We have the following: 

• The "truth", $q\left(x\right)$, which is some distribution that is generating our samples;
• A model, $p\left(x|w\right)$, parametrised by weights $w\in W\subset R^d$, where $W$ is compact;
• A prior over weights, $\phi \left(w\right)$;
• And a dataset of samples $D_n=\{X_1,\dots ,X_n\}$, where each random variable $X_i$ is i.i.d. according to $q\left(x\right)$.

In Bayesian inference, we're trying to infer the posterior distribution of the weights given the data

$p\left(w|D_n\right)=\frac{p\left(D_n|w\right)\phi \left(w\right)}{p\left(D_n\right)}$.

Now, the posterior can be written in exponential form by taking the negative log of the likelihood

$L_n\left(w\right)=-\frac{1}{n}\text{ln}p\left(D_n|w\right)\Rightarrow p\left(D_n|w\right)=e^{-n{L}_n\left(w\right)},$

which gives

$p\left(w|D_n\right)=\frac{\phi \left(w\right)e^{-n{L}_n\left(w\right)}}{{\int }_W\phi \left(w\right)e^{-n{L}_n\left(w\right)}dw}$,

where the model evidence (also called the partition function in physics) is given by

$Z_n=p\left(D_n\right)={\int }_W\phi \left(w\right)e^{-n{L}_n\left(w\right)}dw$.

The expression above is exactly equivalent to the partition function in statistical mechanics

$Z=\int D\varphi e^{-S\left[\varphi \right]},$

where the prior is assumed to be uniform and we've introduced a function called the action $S\left[\varphi \right]$. 

1. ^{^}

   Nanda, N., Chan, L., Lieberum, T., Smith, J., & Steinhardt, J. (2023). Progress measures for grokking via mechanistic interpretability. arXiv preprint arXiv:2301.05217.

2. ^{^}

   Arjovsky, M., Bottou, L., Gulrajani, I., & Lopez-Paz, D. (2019). Invariant risk minimization. arXiv preprint arXiv:1907.02893.

3. ^{^}

   Jaynes, E. T. (1957). Information Theory and Statistical Mechanics. Physical Review, 106(4), 620.
   Jaynes, E. T. (1957). Information Theory and Statistical Mechanics II. Physical Review, 108(2), 171.

4. ^{^}

   I expect to be charged with Moral Realism here, but I don't think that moral realism is necessary for the argument. If you believe there's an equivalence between good epistemic practice and physics (as argued in section 2) then writing an action with a moral field $\varphi \left(x\right)$ is mathematically equivalent to specifying a reward function.