Code

https://github.com/ethansmith2000/TurboSOAP

TLDR

This optimizer keeps the core SOAP idea: rotate gradients into a covariance eigenbasis, perform Adam-style adaptation in that basis, then project the update back.

Standard SOAP periodically recomputes the covariance eigenbasis with eigh or qr subspace iteration. That gives an accurate basis, but it is expensive: eigendecomposition is cubic in the factor dimension. QR is a solid improvement, but is a particularly special torch operation that is less GPU friendly and does not reap the benefits of lower precision computations.

This version adds warm basis refreshes. Instead of solving for a fresh eigenbasis every time, we update the current basis `Q` with a small manifold step, then re-orthogonalize it with a Newton-Schulz polar retraction.

A more extensive description of tangent-step corrections and reothognalization with Newton-schulz was also done for TransportMuon here, explaining a bit of the principles of this kind of method along with visuals.

At a high level:

C = covariance factor

Q = current basis

C @ Q                       # apply covariance to current basis

direction = tangent/Brockett update

Q_new = retract(Q + eta * direction)

The goal is not to perfectly solve the eigensystem every refresh. The goal is to keep Q good enough that Adam-style elementwise adaptation still happens in a mostly decorrelated basis, while avoiding frequent expensive decompositions.

Warm Basis Refreshes

SOAP maintains covariance factors and orthogonal bases for each tensor mode. The usual refresh is:

_, Q = torch.linalg.eigh(C)

This gives eigenvectors of the covariance factor C.

The warm-refresh version treats Q as a point on the Stiefel manifold, the set of orthogonal matrices. We take a small step that improves diagonalization of Q.T @ C @ Q, then retract back to the manifold.

The refresh looks like:

Q -> Q + eta * direction -> orthogonalize

Where eta determines our step size of the correction factor, similar to a gradient step.

There are two useful directions here: a tangent direction and a Brockett-style direction.

Tangent Direction

The tangent direction is the simpler basis-tracking step.

def _soap_tangent_direction(Cq, q):
  diag = (q * Cq).sum(dim=0)
  return Cq - q * diag.unsqueeze(0)

Here Cq = C @ q.

Conceptually, C @ q_i tells us where the covariance wants basis vector q_i to move. The term

q * diag.unsqueeze(0)

removes the component already aligned with each current basis vector. That leaves the part of Cq that actually rotates the basis.

This is a cheap warm-started eigenvector tracking step. If the covariance EMA changes slowly, the old basis is already close, so a small tangent update can keep it useful without a full eigendecomposition.

Brockett Direction

The Brockett direction adds an ordering objective.

def _soap_brockett_projected_direction(Cq, q, weights):

    grad = Cq * weights.unsqueeze(0)

    return grad - q @ _soap_sym(q.T @ grad)

The objective is:

maximize trace(W Q^T C Q)

where W is diagonal with decreasing weights.

This does two things:

• It encourages Q.T @ C @ Q to become diagonal.

• It encourages the columns of Q to be ordered by eigenvalue.

That ordering matters because SOAP’s Adam buffers live in the rotated basis. If the basis columns reorder unpredictably, the optimizer has to keep the moment buffers aligned with the new basis. Brockett-style diagonalization helps make the basis tracking more structured.

The projection term

q @ sym(q.T @ grad)

keeps the update in the tangent space of the orthogonal matrices. So we are not just applying an arbitrary gradient step to Q; we are applying a step compatible with the orthogonality constraint.

Newton-Schulz Retraction

After either tangent or Brockett step, Q + eta * direction is no longer exactly orthogonal.

So we retract it back to the orthogonal manifold.

Instead of QR or eigh, to get back to the Stiefel manifold, this implementation can use Newton-Schulz polar iterations:

Q_new ≈ X (X^T X)^-1/2

where:

X = Q + eta * direction

The practical benefit is that Newton-Schulz is mostly matrix multiplies. That makes it more GPU-friendly than repeatedly launching decomposition kernels, and it can often run in lower precision such as bf16.

Diagnostics, Fallbacks, and Naive Implementation Diffuclties.

Warm tracking can fail if the step is too large or the covariance changes too quickly.

A useful diagnostic is the normalized off-diagonal error:

diag_err = ||offdiag(Q^T C Q)|| / ||Q^T C Q||

If the update reduces this error, the basis is becoming more diagonal. If it increases the error, the step was probably too aggressive.

This lets the optimizer do controller-style refreshes:

try a warm step

measure diagonalization error

accept if it improves, possibly increase step size based on size of error improvement

shrink step size if it does not

optionally fall back to eigh if the basis is too poor

This increases the optimizer time nearly double, which, while thats still not too bad overall, we can amortize by checking this once every N refreshes.

Experiments

LLM Pretraining

• dim=2048

• depth=24

• batch size=32

• sequence=1024

Loss

Loss at 50k steps is about 0.02 higher than baseline, though I suspect this can be improved with tuning of settings around step size and/or normalized steps.

Optimizer Time

• 398ms for tangent mode of TurboSOAP, every 5th refresh (50 steps) costing about 1000ms to recalibrate step size

• ~4000ms for QR

• 13300ms for eigh

Background on SOAP and the Second Moment in Optimizers For Nerds

Most modern choices of deep learning optimizers, in addition to typical momentum of the gradient, often have their own unique form of performing some version of whitening or normalization, typically utilizing a sparse approximation of the second moment of the gradient.

For background, the idea of Moments comes from statistics when describing a distribution, where we consider expected values or averages of different orders, with general form as E[x^n].

• First Moment = E[x] or the mean or average

• Second Moment = E[x^2]

• Third Moment = E[x^3]

• Fourth Moment = E[x^4]

We typically make use of the Uncentered Moments like those above, but we also can have centered ones which have more unique names that are probably more familiar, just to connect it back. The general form is E[(x - E[x])^n]

• First Moment = E[x] or the mean or average (no difference, at least for practical purposes E[x-E[x]] isn’t very exiting)

• Variance = E[(x-E[x])^2]

• Skewness = E[(x-E[x])^3]

• Kurtosis = E[(x-E[x])^4]

Now this is the case for when we consider 1-Dimensional distributions. When we work in higher dimensions, the shape of the metric grows with each order. For a space of dimension N

• First = N

• Second = NxN | uncentered E[xx^T] | centered = E[(x-e[x])(x-e[x])^T]

I’ve bolded the E[xx^T] as this is the object we’re typically estimating, though usually only sparse approximations of it.

Why?

When we have a parameter of size m x n, we know its gradient and momentum are the same size. Though when we want to estimate the second moment, we now need to consider every dimension multiplied with every other dimension. For the purposes of optimization, our m x n matrix is a somewhat symbolic structure, there are times it makes more sense to think of it as a vector of dimension m*n.

So, when we think about our second moment, its full size, to capture the interaction of every dimension with every other dimension is mn x mn. For at 1024x1024 matrix parameter, this is 1024^2 x 1024^2 which is 1 million by 1 million.

Obviously that’s unreasonable for us, so we typically aim for sparse forms of it. Probably the two most common are:

• Diagonal

  • Only considers variance (each dimension’s interaction with itself, no cross interactions), the square root of that is what we normalize by.

  • The diagonal of the mn x mn matrix is just mn in size

  • This is what Adam does for normalization via its denominator. This is somewhat like the idea of a z-score but instead of typical (x-mean)/std we have moving averages at different speeds comprising something like ema_mean@0.9/ema_uncentered_std@0.999

• Rank-1 Kronecker-like or similar factorization

  • Here we more explicitly consider the 2D structure of our gradient

  • For gradient g, in its 2d shape, we are often estimating both E[ggt] and E[gtg], each symmetric matrices representing covariance, which are then used to form both “left” and “right” preconditioners.

  • Similar to the Adam case, of applying v^(-1/2) (bringing us back to x-space from x^2-space per-se) to get our normalization factor, one form of normalization is to raise each of the matrices to -1/2 power.

    • In some cases, this is -1/4 power, since we have both the left and right to apply.

  • The resulting normalization matrix we get is often referred to as a preconditioner

  • One way we could feasibly apply our preconditioners is to do

    • P = (P_left[:, None, :, None] * P_right[None, :, None, :]).flatten(mn,mn)

    • pg = g.flatten() @ P

    • This should remind you of how we can take the outer product of two 1d vectors to yield a 2d matrix, which can be done by matmul or
      A[:,None] * B[None,:]. This is the same idea generalized to taking two 2d matrices outer product to form a 4d matrix, and flattening it.

  • Recall that the outer product is effectively every dimension matched with every other dimension.

  • Overall, this is an attractive structure as it “acknowledges” the 2d structure of our network weights

  • Some examples include Shampoo and PSGD

Enter SOAP

Now SOAP Optimizer (ShampoO with Adam in the Preconditioner’s eigenbasis) is one particularly interesting and powerful optimizer. Its something of a hybrid of both styles, and can be quite powerful.

SOAP involves a normalization of the gradient/momentum in the eigenbasis coordinates of the gradient covariance in a Adam-like fashion. .

From our left and right (Exponential Moving Averages) EMAs, E[ggt] and E[gtg], we perform eigendecompositions (as A=QΛQ⊤ ) to obtain an orthogonal eigenbasis.

The idea of normalizing coordinates of orthogonal direction may remind you of Muon

Muon uses Newton-Schulz to perform a 0th power on the singular values of the gradient/momentum matrix itself (approximately aiming for uniform 1s of the singular values) but does this without explicitly performing SVD to obtain the singular values. In other words, one could imagine doing either SVD to obtain USVT and raising S to the 0th power to make S all ones (or just directly tossing it out) and fusing UVT. Meanwhile the output of Muon directly is an approximation that converges to UVT

Muon is approximately taking

M = U S V^T

and applying:

M -> U V^T

which is the polar factor of M. Equivalently, it replaces the singular values by ones, up to rank/rectangular-shape details.

So SOAP normalizes coordinates after rotating into an estimated covariance eigenbasis. Muon normalizes the singular spectrum of the update matrix itself.

The distinction is:

SOAP: covariance eigenbasis + elementwise Adam normalization in that basis

Muon: polar/orthogonalized update direction of the gradient or momentum matrix

For an m x n gradient matrix, Muon’s singular values are min(m, n) values from the matrix being updated.

SOAP’s covariance eigenvalues come from the row/column covariance factors: m eigenvalues for G G^T and n eigenvalues for G^T G for total of m x n, effectively per-coordinate (mn) of a factorized eigenbasis that is mn x mn.

Difficulties with SOAP at scale

Memory

There are difficulties with Matrix-based optimizers in general at scale but SOAP may be particularly difficult due to a large volume of persistent states.

If we think about an n x m parameter, we can talk about memory cost in terms of that. For example, Adam maintains a momentum and second moment term of the same dimensions. So its net memory usage is 2 x n x m.

SOAP on the other hand has the EMA of the gradient covariance on both sides, the preconditioner on both sides, and the momentum + diagonal second moment term. The net price you pay is 2 x n x n + 2 x m x m + 2 x m x n. This especially can bite on large rectangular matrices like MLP’s.

Distributed/Sharded Training

FSDP and other parallel methods have generally been pretty easy for Adam as all of its operations have been elementwise, considering each dimension in isolation. So formulations that have flattened parameters or sharded them like DTensor have had little issue.

For matrix-level optimizers, it’s been trickier as we have two conflicting requirements.

Our states and gradients, by default, are split between devices. However, we need to materialize the full matrix, gathering all shards to perform our matrix operations typically. This requires all devices to sync to communicate their shards.

However, at the same time, similar to how when we maintain the shard form with Adam, we’d also really like to split the responsibility of updating parameters across devices. Effectively, it would be desirable to evenly distribute our parameters across devices (or by a heuristic of expected computation time) and let each operate asynchronously.

So we’re fighting a need for sync against also wanting to be async, and the naive form can be pretty brutal. There are improvements but it can be tricky.

This work addresses the speed component, though the two points above still pose difficulty at scale, a common problem with matrix-level optimizers.

Code

https://github.com/ethansmith2000/TransportMuon

Premise/Question

The original Muon, at optimization step t, applies the Newton-Schulz 5th-order iteration to the gradient momentum matrix M_t to produce an update Q_t. Q_t approximates the polar factor of M_t:

Since momentum uses a well-sized decay factor, e.g. beta = 0.95, consecutive momentum matrices M_t and M_{t+1} usually change slowly. Should the target polar factor U_t V_t^T also change slowly between optimizer steps? If so, can we avoid recomputing Q_{t+1} from scratch with several Newton-Schulz iterations, and instead cheaply update or correct the previous polar factor Q_t toward the new target associated with M_{t+1}?

High-Level Idea

Making use of Q_t and M_{t+1}, we’d like to reach the point of Q_{t+1}.

A difficulty is that Q lives on the Stiefel Manifold, a manifold of m x n matrices where m >= n and the columns are orthonormal. This both influences how we consider measuring error/distance and creating an update.

Generally there are two schools of thought for this problem

• Intrinsic/Geodesic - the update itself stays directly on manifold

• Extrinsic/Retraction-based - we take a step off in ambient euclidean space/tangent space and then retract back onto the manifold

The first requires solving the exponential map, which can be expensive. The second has some nice tools we can make use of, and we already know an efficient retraction method: Newton-Schulz

I find the methods are nicely visualized when considering a simple manifold like the 2D unit-norm hypersphere (though i generally prefer to call them circles).

The TLDR

Take a tangent step off manifold, aim to use only one newton step to retract back onto Stiefel manifold as faithfully as possible to the UVT representing the new momentum term, or the polar factor we would get from performing a full newton-schulz loop on the new momentum from scratch.

Every N optimization steps, perform an “anchoring” where we take a fresh momentum term and converge as closely as possible to singular values of 1. For this, I used 5 steps of the typical Newton-schulz-quintic tuned by Keller Jordan. While this is a strong tuning for converging small singular values quickly, it notably does not converge singular values closely to 1.0. So, I follow up with 2x steps of cubic newton-schulz with a converging tuning.

Then for next N optimization steps, do 1 tangent update + 1 step NS to retract in place of full newton-schulz 5 steps.

For all Newton-schulz cases except for when we run the NS5 from scratch, instead of Frobenius norm, it’s been preferable to cap by an estimated top singular value.

First: The Update

Given the previous cached update Q_t and the new momentum M_{t+1}, we can measure misalignment via the skew-symmetric part of:

If Q_t is already close to the new polar factor, then R should be nearly symmetric.

We know this by this identity

Q_{t+1} M_{t+1}^T = U V^T (U S V^T)^T
                  = U V^T V S U^T
                  = U S U^T

The alignment error is the skew-symmetric component, effectively how mirrored positions differ, the quantified violations of symmetry.

Or elementwise

So if we can reduce the skew, this would signal we’re approaching our desired target.

We can frame our problem like this:

Aside on this formulation:

Let A is skew-symmetric (A^T = -A; A + A^T = 0)

This choice of A naturally lets AQt lie in the tangent space. Taking a step in this direction:

• First order error cancels, but we do have a second order error.

• AQ_t is a “rotation-like” update to the rows of Q_t. Because A is skew-symmetric, it mixes rows in opposite directions: if row i moves toward row j, row j moves away from row i by the matching amount, preserving orthonormality for the infintesimally small step

• Though we use a finite Euler step

• This induces the second order error on both our orthogonality and inflates our singular values

Let’s imagine a small step of our update

And then expand:

Q_eps Q_eps^T
= Q_t Q_t^T
  + eps (A Q_t Q_t^T + Q_t Q_t^T A^T)
  + eps^2 (A Q_t Q_t^T A^T)

If Q_t Q_t^T = I, then the first-order term becomes:

Which cancels out by the skew-symmetric identity.

Looking at the second order term, it reduces to

-A^2 is positive semidefinite, so this will inflate singular values and screw up our orthogonality, hence the need for the retraction step. Remember that this error exclusively increases singular values; that will be a valuable property we can exploit in our design.

Back to how we get A

We said that the update direction should reduce the skew of R (QMt). So starting from our original problem framing, of wanting to obtain the correct polar factor (Q), let’s get to R by multiplying each term by M_{t+1}

From:

to:

simplified:

and our target point is such that

Solving this exactly is too expensive. Instead we can use a Jacobi-style diagonal approximation. For an off-diagonal pair (i, j), approximate the effect of A using mostly the diagonal entries of R.

(AR)_ij ≈ A_ij R_jj
(AR)_ji ≈ A_ji R_ii = -A_ij R_ii

Therefore:

skew(AR)_ij = 0.5 * ((AR)_ij - (AR)_ji)
            ≈ 0.5 * A_ij * (R_ii + R_jj)

to cancel the skew error:

skew(R)_ij + 0.5 * A_ij * (R_ii + R_jj) ≈ 0

Finally:

A_ij ≈ -2 * skew(R)_ij / (R_ii + R_jj)

The negative sign moves against the skew error and the factor of 2 cancels the 0.5 in the skew definition (which in code we just omit the 0.5).

As in code (0.5 on skew included for clarity):

R = q_prev @ work_matrix.mT
skew = 0.5 * (R - R.mT)

diag = torch.diagonal(R)
denom = diag.unsqueeze(1) + diag.unsqueeze(0)
floor = jacobi_eps * (diag.abs().mean() + 1e-8)
denom_abs = denom.abs().clamp_min(floor)
denom = torch.copysign(denom_abs, denom)

A = -2.0 * eta * skew * denom.reciprocal()
q_tilde = q_prev + A @ q_prev

Second: The Retraction

The canonical Newton-Schulz used in Muon converges like this (green 1step + red 5step lines)

This is a very deliberate design choice. The singular value distribution we have in our gradient momentum is a black box, and its possible (and even likely) some singular values may be quite small.

So this variant is strong is moving small singular values quickly.

But it does not converge monotonically to singular values of 1.0, instead oscillating around [0.7, 1.3].

In practice, this has been fairly sufficient, but because we want to take lots of tangent steps from a single polar factor, we might want our initial starting point to be on manifold as much as possible.

Keller shares another set of coefficients (pictured here in the solid lines) as a base example. I started with this one, but tried exploring some others balancing between things like

• Convergence speed of small values

• Convergence point

  • Fixed or oscillating

• Maximum values before we risk explosions

A general rule for convergence to 1.0 is to have: a - b + c = 1.0

This is another

As mentioned previously, when we do our tangent update, we exclusively increase our singular values.

From there, we can do a singular value upper bound estimate and normalize by that, which happens to use the gram matrix, making it a small detour away from our typical NS path.

What number we use for our cap depends on bit on where each form of NS starts to blow up.

I haven’t experimented with it as much, but personally could imagine benefit with one of the variants that has fast single step growth for smaller values, even if it overshoots a bit, we could then keep capping at 1.0.

Full Thing in Pseudocode

Q_cur = None # Polar Factor (the update)
Q_prev

For i in train_steps:
    g = model_loss()
    M = update_momentum(g)
    if i == 0 or i % anchor_steps == 0 or i < warmup_steps:
      Q_unstable = NS5(M)
      Q_cur = NS_cubic_2_step(Q_unstable)
    else:
      Q_tilde = tangent_step(Q_prev, M)
      Q_tilde = clamp_by_largest_sv(Q_tilde)
      Q_cur = NS_cubic_1_step(Q_tilde)
    update_model(Q_cur)

Experiment1

• 1024 dim

• 12 layers

• 1024 seqlen

• 32 batch

• Muon default settings from original repo, 0.02 lr, 0.95 momentum

Initially off to a rough start once we switch from full NS warmup to transport method. I imagine I did not warmup for long enough

Crosses below baseline around 25k-ish

I suspect it possibly performs a bit better given that for our transport steps, and anchor, we approach very close to singular values of 1.0

It was a bit hard to display, so I disabled the other runs for the optimizer time. In ms, we have 85ms every 10th step, with other steps being about 11ms shorter.

• ~32 seconds gained on anchor step

• 11x9=99ms seconds saved on transport steps

Experiment2

Same as above but now 2048 dim

This time I decided to warmup for 2500 steps instead of 500

Base Muon takes the lead during the warmup phase, but right around the start of our anchor-transport loop, TransportMuon takes a significant dip. I was not expecting this. I figured it could be related to where singular values converge to, inflating learning rate, but given that I cap at 1.25 + retract, 1 step of NS is expected to give us 1.046 at the max.

Around 40k, pretty similar delta in loss.

The optimizer speed on the other hand is a pretty comfortable improvement.

• ~27.5ms gained on anchor step

• ~48.5*9=436ms saved on transport steps total

Related Work

• Jordan, Jin, Boza, You, Cesista, Newhouse, & Bernstein. Muon: An optimizer for hidden layers in neural networks.

• Carlson, Collins, Hsieh, Carin, & Cevher. Stochastic Spectral Descent for Restricted Boltzmann Machines.

• Tuddenham, Prügel-Bennett, & Hare. Orthogonalising gradients to speed up neural network optimisation.

• Bernstein & Newhouse. Old Optimizer, New Norm: An Anthology.

• Bernstein. Deriving Muon.

• Bernstein and the Thinking Machines team. Modular Manifolds.

• Cesista. Heuristic Solutions for Steepest Descent on the Stiefel Manifold.

• Cesista. Steepest Descent on Finsler-Structured (Matrix) Manifolds.

• Su. Stiefel Muon.

• Tilde Research. Gram-Space Manifold Muon.

• Manifold Constrained Steepest Descent (MCSD / SPEL).

• Kong, Wang, & Tao. Momentum Stiefel Optimizer, with Applications to Suitably-Orthogonal Attention, and Optimal Transport.

• Edelman, Arias, & Smith. The Geometry of Algorithms with Orthogonality Constraints.

• Absil, Mahony, & Sepulchre. Optimization Algorithms on Matrix Manifolds.

• Bécigneul & Ganea. Riemannian Adaptive Optimization Methods.

• Kasai, Jawanpuria, & Mishra. Riemannian Adaptive Stochastic Gradient Algorithms on Matrix Manifolds.

• Li, Fuxin, & Todorovic. Efficient Riemannian Optimization on the Stiefel Manifold via the Cayley Transform.

In our pursuit to perfect neural networks, we often look to how humans learn for reference, which has had varying degrees of success. Whether or not these inspirations lead to performant models, the research itself, both positive and negative results, is interesting, and I like to think it only adds to our knowledge compendium.

One space of human-inspired learning I am particularly interested in is how humans learn in social settings. The human, rarely if ever, learns alone. Unlike our typical methods, we don’t chug through data in isolation. Imagine, for a second, if that were the case. We’d go through life seeing, reading, feeling, and hearing our environments. We’ll also allow watching others behave as neural networks take in data like this as well.

However, we’d remove many places for learning we have, such as debates with our peers trying to enforce our world models upon each other, dealing with cognitive dissonance when we experience mismatches between reality and our notions of the world, assembling the perspectives of those in proximity with us, balancing others’ opinions to help gauge our own, and so forth. These behaviors, I would argue, do not have well-defined analogues in neural network learning.

One could argue teacher-student learning could constitute an example of social learning, though it feels more like distillation and a single-channeled knowledge transfer. Namely, I would be interested in cases where a neural network both learns from its own data stream but is also forced to weigh its developing world model against others with distinctly different ideas. Methods like Bayesian learning do resemble the idea of balancing our priors with new incoming information, though we somewhat predefine the weights of how we balance this dynamic.

How might a neural network develop its own agency to decide what information is not worth adding to its knowledge base, writing it off as noise or invalid, and what is worth learning? I believe we’d call this bias.

You’ll hear many people refer to neural networks as biased. This couldn’t be further from the truth. Our models have slight biases towards solutions that generalize well and are “smooth” by dynamics of gradient descent and some other magic we still haven’t fully understood yet. Though, overall, they are very close to unbiased estimators of the distribution of observed data. They approximate mirrors of the reality they were shown. When we accuse neural networks of acting in biased ways, what’s really happening is that its training set was sampled in a biased fashion, and our model just absorbed that bias without its own bias. For instance, when neural networks generate images of professions with a noticeable demographic bias, you’ll typically find a very similar bias performing a search on Google. Thus, it is the bias of humans selecting how we build our digital world that results in a collection of images on the internet that are not representative of reality that the neural network later digests. The distribution of images on the internet is in no way representative of how often they occur in reality, which not only hurts the chance of neural networks learning accurate depictions but also further biases ourselves as we see a distorted picture of reality.

So I hope this example drives home the point that neural networks effectively learn an unbiased reflection of their dataset (indicated by low cross entropy or FID), while humans are biased creatures. Unbiasedness seems like a desirable trait to have for fairness and accuracy in our neural networks. Though, I believe human biases exist for a reason outside of discrimination and other maladaptive behaviors. Our biases serve us to select for information that we consider valuable or coherent with our existing notions and reject noise or dissonance. This kind of intuition feels nontrivial to construct in current neural networks.

Somehow networks would need to develop their own sense of how to weigh incoming information by utility, preference, and consonance, not just by surprisal factor, predefined loss clipping, and a KL distance to another model. Perhaps reinforcement learning could account for some of this, but I’d be interested in how we could better understand this phenomenon and design training methodologies that replicate the ways we acquire knowledge through socialization and use biases to filter what influences our world models.

I have no idea how this should look, though I have a few mental models of things that could be fun to explore for social learning examples despite not having a great outlook for profound results.

Socially Cooperative Learning

A method of training where a fleet of models each see different shards of a dataset, but each model also depends on the others for filling out the gaps of what is not available in their shards.

A possible scenario:

1. Initialize N models of the same type (using the same initialization?)

2. Split the dataset into N evenly sized shards

3. At every training step

   1. each model trains on a batch from its own shard and updates its weights

   2. each model takes a step towards all of the other models by linear interpolation

Another possible scenario:

1. Initialize N models of the same type (using the same initialization?)

2. Split the dataset into N+1 evenly sized shards, one shard will be communal

3. At every training step

   1. each model trains on a batch from its own shard and updates its weights

   2. all models make a prediction on a batch from the communal shard

   3. predictions are averaged, and this becomes the target for each model to train on as opposed to using ground truth labels of the communal set.

Cons:

• This is quite expensive both in computation, loss of parallelism, and memory for storing multiple models. Would probably only try this with SSL or classifier models.

• It’s uncertain how much of an advantage is here compared to EMA training or post-training model soups/merging

• More hyperparameters to decide how large of a step models should take towards each other, although there may be a heuristic that fits here, like weighing the step size by distance.

The hope:

• Improved generalization capabilities as indicated by performance on the validation set and reaching flatter minima.

• It could be interesting for decentralized/async training options where model weights exhibit a controlled divergence across devices instead of being identical copies and having the update step with other models be less frequent.

• In the end, the models would be used as an ensemble, weighing all of their predictions together or averaging all model weights into a single model.

Socially Adversarial Learning

A method of training where 2 generative models face off against each other, attempting to “persuade” the opponent model towards becoming more like itself. In short, a model aims to generate the optimal dataset such that when trained on, the opponent model “agrees” more with it.

A possible scenario:

1. Obtain 2 fine-tuned LLMs, each fine-tuned from the same base model on different sets. It’s likely important that the models have meaningful distance from each other, i.e., not equivalent solutions that are permutations of each other.

2. Obtain a dataset of answerable questions or prompts for the LLMs to respond to

3. Each LLM writes an answer to the question; the output is fed to the opponent to train on.

4. The reward each model gets is the cosine similarity between the opponents update direction and the difference between the two LLM weights, scaled by the update’s norm

5. To regularize, each LLM also takes a small step towards its original weights (or KL-style regularization) to prevent collapse and stay somewhat anchored to its original self.

Cons:

1. This kind of RL would likely have high variance and be difficult to train

2. Outputs from models could end up appearing like adversarial examples rather than being meaningful. In other words, generating text that is persuasive with respect to affecting another model’s weights might not necessarily make a model persuasive.

The hope:

1. Not as sure here! There is some hope that models can use adversarial games for continual self-improvement, but I’m not quite sure what to expect here.

2. Another example may be more purely in the data space rather than fussing about with model weights. For instance, two LLMs debate each other, a winner is somehow decided, and this becomes the reward and penalty to improve one’s debating abilities.

Doing this recurrently until convergence yields what I’d like to call a consensus algorithm.

A consensus algorithm smells a lot like unsupervised clustering algorithms in that they identify groupings of points, but with the added component that the values of points themselves are modified at each iteration to move closer to their neighbors until all points reach a “consensus” or an agreed-upon value they all take on.

To make this more concrete, we can say more on how it differs from something like K-means clustering.

K-Means

• Initialize centroids in the same space as the data

• Move these centroids around in order to have minimal distance from the data points in their cluster

• Reassign data points to centroids based on how close they are to each centroid.

Consensus algorithm

• Data points themselves serve as the centroids

  • No predefined number of clusters: potential to have as many clusters as there are data points. Hyperparameters used to define the update step may influence how many points we converge upon

• Instead of updating the initialized centroids at each step, we update the data points themselves to take a step in a direction that considers all other data points, with the directions corresponding to their closest neighbors having the most importance

Parameterless attention is one possible consensus algorithm. I originally discovered this use case when trying to design a color clustering loss that could be differentiable for backpropagation and benefit from GPU acceleration. To note, despite it being differentiable, it did not work very well as a loss.

Assuming you know how attention works, and specifically the QKNorm variant, it works as follows:

• For each data point:

  • Calculate vector of cosine similarities with all other points

  • Apply temperature scaling, a hyperparameter that decides how strongly to prioritize nearer points over farther away points

  • Apply softmax over this vector, giving us the mixing coefficients or weights.

  • Take a weighted sum of all the data points scaled by the mixing coefficients

  • Perform the update: a linear interpolation between the existing data points and the output of the attention operation we just did.

This should probably ring a bell for attention methods used in graph neural networks.

It looks like this

We can use it to perform segmentation of an image in color space

Or, by initializing random noise centroids instead, and performing something closer to “Perceiver Attention“, we can perform something closer to a “soft” K-means algorithm.

We can also perform consensus in spatial domains.

I initially tried using the same hyperparameters as before but found that the points did not move at all. With the scale as high as it was, it’s possible that nearly all the weight was put on self and immediate neighbors.

scale=100.0

step_size=0.85

at smaller scales, we just collapse to the mean, which is probably pretty reasonable behavior for single gaussian, which really is just a single cluster of points.

scale=1.0

step_size=0.85

This time around, I initialized 3 gaussians with means (1,1), (0,-1) and (-1,1). This appears to reach consensus with 3 clusters as we may expect.

scale=10.0

step_size=0.85

The code I used for this is here:

https://github.com/ethansmith2000/attention-cluster

I felt compelled to make this post after seeing yet another reinforcement learning paper for diffusion models that does spectacularly in terms of being able to fit to the reward function, but the actual results look terrible, collapse to a single style, and are otherwise kitsch.

This is not at all a criticism of the author’s work. They proposed a new method of reinforcement learning, and it successfully fit the objective quite well. But the images all have the same tired, fake HDR, airbrushed look that make them all painful to look at.

I’m reminded of the diagram when OpenAI intended to show that DALLE3 had improved over their previous DALLE2, but the newer outputs were just so much worse.

The problem is that the reward function itself, I think, sucks.

In this post, I want to discuss the nature of human preferences and why standard reward models, which learn the expectation or average over many human labels, might be doing us a disservice.

The Nature of Human Preferences

The kinds of things humans like differ wildly from person to person. As for what makes a pleasing image, some people may prefer a piece by Van Gogh, others may prefer anime drawings, and others may prefer abstract art.

When collecting ratings of these things across many people, we might get wildly different answers.

Inevitably, because we all have our own subjective preferences, there is a lot of variance and inconsistency in human ratings. One person’s trash is another’s treasure.

We hope that by getting many unbiased ratings, we can “average out” this noise and reach the common denominator of what makes an image aesthetically pleasing, regardless of the exact kind of content.

In other words, while some may prefer oil paintings and others may prefer 3D-rendered art, what are the universal components that are pleasing to all? Maybe certain colors are universally enjoyed on average by a certain population. Maybe there are certain compositions or uses of light that are generally preferred.

However, I think if there is any preference that persists on average across such a large population, it is something so primitive and vague that it entirely misses the individual nuances of what people actually love.

I say this because many aesthetic preferences are mutually exclusive. Let’s imagine a toy example of this.

Let’s say half of a population enjoys dark, brooding art, and the other appreciates the opposite: light and happy scenes. We then collect a dataset composed of these two art styles, again half and half. The dark crowd always rates the dark art highly and the light art negatively, and vice versa.

For our model, this is unfortunately not an informative trait it can use to make predictions. If half of the labels over dark art are rated highly and the other half are rated lowly, the minimum error prediction is to predict a value right in the middle, effectively making this trait invariant to its prediction.

Instead, the model will need to depend on other features that on average yield higher ratings. Like, let’s say pictures of people are rated more highly on average by both groups; the low ratings are slightly higher, and the high ratings are also slightly higher.

I would worry that at scale, over many different ratings, a lot of the nuance of what makes a certain image preferable is canceled out, and we are left only chasing a mean that embodies a bunch of “safe” traits that are meant to be universally receptive but appeal specifically to no one!

The space of human reward functions is extremely variant and often contradictory. Here, the floor axes represent different kinds of images, and the vertical axis shows how highly a certain person may rate them. For a model to estimate such noisy variant scores from many people with the least error, it would learn the mean preference, which, notice, is not rated highly by anyone here. We’re trying to fit a single reward function to many highly incompatible reward functions that each exist in our data in conflicting ways, which manifests as highly variant, inconsistent datapoints, making it difficult to fit to.

In neural network modeling, we have explainable variance and variance from noise. Explainable variance is something we can exploit in learning, variance from noise is just something we have to deal with. If we trained a classifier on cats and dogs, we’d hope that the change in visual features corresponds nicely to the change in labels. However, if the labels are unreliable and noisy, like if we just randomly labeled each image, the model would likely converge to outputting p=0.5 for every image, as this would be the safest bet it can make with no pattern it can reliably depend on.

The variance we have in labels and scores here, while it is sampled in a meaningful way, may manifest as noise due to trying to learn from partial subsets of many exclusively different functions.

When our reward dataset is more objective, say, successfully solving math problems, rendering text that can be read correctly by an OCR, or the compressibility of an image, I think we can get away with it though.

A finding that might support the claim is that the LAION aesthetic predictor v2 was trained on a dataset of images labeled 0-10. Though in practice, we observe in generating predictions over millions of images, something like 99% of the data falls between 4-6, suggesting that the model learns a very conservative fit to the data due to the noisy labels. Granted, I also don’t know the nature of the labeling process.

Another way I like to visualize this problem is to imagine doing time series predictions.

Let’s imagine the pattern in blue occurs many times in the dataset. Half of the time it continues steeply upward, and the other half of the time is continues steeply downward.

Just predicting the mean tells us a terrible story about what possible continuations might look like.

If we predict mean and variance, we at least have a sense of spread, but still the majority of probability mass is on the center, which is inaccurate.

Once we get to Gaussian mixture model territory, we can respect the multiple clusters of outcomes.

Now, if we can accept that the typical strategy of reward modeling ignores a lot of nuance and focuses on the universal component instead of respecting clusters of human preferences, then we can consider some solutions. First we’ll discuss solutions for creating the actual reward model. We also need to figure out how to use that reward model to train an image model.

Training more nuanced reward models

1) Gaussian Mixture models

Noting that our data is composed of many multimodal distributions we’d like to learn, we can try using a model that is capable of fitting to such distributions. Namely, the Gaussian mixture model is composed of many Gaussians.

Let’s visualize our light and dark example. Again we have very split distributions of scores

With a mixture of Gaussians, we can learn to divide and conquer. One Gaussian can take on modeling the light crowd, and the other can model the dark crowd. Because our data is split half and half, we’d likely also see that the mixing weights are [0.5, 0.5], which, if we were to mix, would have us right back where we started. But now, we have effectively factorized the view, we can assess looking at each Gaussian independently.

One issue is that the number of Gaussians we choose to model with is an integer value. We don’t know how many Gaussians we should be using, which is a common issue in unsupervised clustering problems, i.e., how many means to use in k-means. However, all cases might still be more useful than the base case of a single Gaussian, and we can choose how much granularity we want.

2) Conditioning on the rater

If we have information about our raters, we could provide that as a condition to the model, factorizing the full distribution of ratings into many conditional Gaussians, where each one models a specific crowd. If we have any metadata as well about scores, i.e., users from Australia typically rate art from the 1970s with X score, these factors could also be engineered into a condition. We’re now modeling p(rater|x) for across many raters but still able to get p(x) through marginalization

3) Z-vector or random condition.

Similar to the spirit of VAEs, we may be able to inject a randomly sampled z-vector in one of the deeper parts of the network, and get a different rating per z-vector we provide

Reinforcement Learning with our nuanced reward model

A problem now is that in all of these cases, we are outputting multiple scores. The GMM outputs as many scores as we have Gaussians, conditioning on the rater gives us a possible score per rater condition, and there’s a similar scenario for the Z-vector case.

We could aim to maximize all scores, but then we’d be back at our original problem of getting an average output that is universally pleasing to all raters.

Instead, I think we should condition the generative model on the specific rater/vector and accordingly reward it as such. We can think of this like task embeddings, i.e., “make an output that is pleasing to rater 5.” In other words, we now have many, many reward models and many input conditions describing the task scenario. It should be equivariant, meaning as we swap the nature of evaluation, we also swap the nature of the instruction/task.

This expands our training space substantially, since we can train on a single image but also have many different scenarios to evaluate it on.

Something I’ve been thinking about a lot lately is how humans handle noisy recognition. Maybe you recognize the image above, if not you probably haven’t seen the original image. The fact is, many people do astonishingly recognize the person in the image, despite the fact that the image has been substantially downsampled and destroyed. However, the message/signal, the identity of the person, remains intact for us. How?

Meanwhile, AI visual models tend to fail hard at these perturbed recognition tasks, being vulnerable to even relatively mild amounts of image distortion.

The fact that current AI models can’t recognize this feels like a missing piece. Though another part of me wonders if we should even want AI models to be able to recognize these edge cases.

Not for reasons around robots beating Captcha’s or anything like that. More, is it a desirable trait for vision that we look at a smear of pixels and hallucinate something there that isn’t exactly present?

We can look at things like ink blots and clouds and see entire stories within them, and while this can serve as an interesting microscope into the psyche, is it a beneficial trait to have to project our minds onto reality?

Here, I want to discuss the status quo of how AI currently “sees”, how we might strategize an approach that allows for seeing obscured messages through “controlled hallucinations”, and whether this might be desirable.

How AI Currently Sees

Computer vision, in most familiar forms, operates in a bottom-up fashion. An incoming visual signal discretized onto a grid of pixels arrives at the input layer of a neural network and is propagated all the way through to the end to be processed into an semantic representation. In short, image -> conceptual representation. This flow is unidirectional.

An alternative form of flow is top-down processing, where we proceed in the reverse direction, beginning from the idea and mapping back to the visual stimuli. In short conceptual representation -> image. This type of processing is typically reserved for generative tasks, where we may provide a text caption such as “a cat” to guide a model towards generating a picture of a cat.

There is some nuance we’ll uncover as we go on, but these modes often appear separately from one another, and for tasks such as image classification, segmentation, encoding, and visual language modeling, it typically manifests as somewhat exclusively bottom-up.

Vision, at minimum, requires bottom-up processing. We need some means of taking in a signal and further processing it. However, it is possible to have have visual processing that makes use of both forms.

How Humans See

This section is summarizes this video and also borrows visuals from it.

When perceiving, humans make use of a mixture of top-down and bottom-up processing.

Higher cognitive areas constantly generate models of what we expect to perceive based on prior knowledge, context, and recent experience. These predictions flow downward through the neural hierarchy.

Meanwhile, raw sensory data flows upward from our sensory organs. Importantly, what moves up isn’t all sensory data, but primarily the prediction errors - the differences between what was predicted and what was actually sensed.

The brain assigns different weights to prediction errors based on their estimated reliability. In clear conditions, sensory input gets higher weighting; in ambiguous situations, prior expectations may dominate.

The two streams attempt to align with each other and meet in the middle. We hold “hypotheses” about what might be in front of us, and send them downward. Meanwhile, the view of the object in front of us enters through the eyes and is sent upward. In other words, the final result of what we interpret is a mixture of what we guess and what actually enters the visual stream.

This is an efficient design that avoids computational waste by leveraging existing knowledge.

If the two streams agree early on, then there is no need for further processing, our hypothesis was correct. If not, the gates are opened for the visual signal to propagate further, and we adjust our hypotheses with the new incoming evidence. If you stare at the same thing for a while, there is no purpose in repeatedly processing the entire visual field. This becomes our new “default”, and we then only need to process the changes or residuals from this new default. We depend on expectations, and only process unexpected information, which then becomes part of our expectation. The entirety of this scheme is known as Predictive Coding.

Additionally, the top-down stream allows for generating plausible predictions when information is incomplete. For instance, if we see a cat walking out from around a wall, but can only see its head, based on the entirety of our past experiences, we make the reasonable leap of faith that this head is attached to a cat-like body with the same fur coloration as well.

The top-down stream, in a sense, has the final say in this operation. We constantly generate “guesses” and then seek evidence to either confirm or reject our guess. Though it is also responsible for adjusting its beliefs when encountering errors or mismatches with new information entering our sight, aiming to minimize disagreement or “Free Energy”

The top-down component is also quite important due to how little of the world the eyes actually take in. Artist renderings to be taken with grain of salt.

There is much in our vision that we see that isn’t actually present, and some things that are present but only obstructive, like our blind spot. Our top-down processing can learn to ignore the constant parts and fill in the gaps of the parts we need to know.

This assumptive mechanism is also the same reason optical illusions can fool our visual system, like how an inside out mask still looks right side up given how used to we are seeing faces in this manner. We have learned a number of predictive context clues depending on lighting, shadows and contiguity of edges that can sometimes fail us when applied in edge cases. It is also the reason we may see things in the dark that aren’t there-- our top-down processing is still generating guesses, but it has little to no sensory input to keep it in check, causing it to dominate the resulting interpretation. Psychoactive drugs which alter our perception, can also be explained as interfering with the balance between top-down and bottom-up, as well skewing the actual mechanism of each.

Perception that works by mixing our expectations (priors) and raw senses (evidence) to produce the final result (posterior) is considered a Bayesian approach, a statistical framework that takes into account not just what we learn but all of what we already know.

We know at least this method of perception is desirable from an efficiency standpoint, but here, we’ll imagine AI has infinite compute. What is in question is whether the strategy of cycling guesses with evidence until agreement is reached is a trait we want AI to have.

Dreaming is also an interesting case here. In a dream you both construct an experience, and experience it as well.

Thus far, it appears that it can be useful for generalization to new unseen scenarios, but optical illusions reveal a few failure modes. Though ultimately, whether the benefits outweigh the cons is up for debate.

These controlled hallucinations are incredibly useful when they are right, but are problematic when wrong. So if we do opt for this route, it seems essential to understand how to build in assumptions that are correct most often and can adapt readily to failures.

AI that performs both bottom-up and top-down

There are a few cases of neural networks, I would argue, that do make use of the dual-stream approach.

While a pure one-step generative model like normalizing flow models or GANs do not appear to be exercising bottom-up, multi-step methods like diffusion do.

Diffusion model generation begins with a canvas of pure noise, which we can analogize to a lack of signal, or a scrambled signal, like looking at a cloud or seeing in the dark.

At the very first step, without much context to work with, our guess at what the final image looks like is primarily top-down.

From our guess, we then slightly update or refine our canvas to add a bit of context to it, now treating it as signal we will then read bottom-up in the next step.

As we continue removing noise from the canvas the data is revealed, lessening the role of top-down processing and increasing the role of bottom-up.

Typically when a diffusion generation is visualized, we are observing the current state of the noisy canvas. However here is a visualization of what our guess is of the final denoised output at every step. In early steps, this changes significantly. Our context is low, so our hallucinated end result is in flux. But as we add more context, it starts to solidify as agreement is reached.

A key difference here is that we are allowed to update the data with our guesses. Normally in human perception, your guesses need to adapt to the static unchanging phenomena. However, here the guesses actually change reality and modify what is received bottom-up on the next step. The model takes a look, add its own hallucination on top of it, rinse and repeat until the hypothesis and the visual signal become the same.

Another example of this is Google’s DeepDream. DeepDream takes a model that is purely bottom-up but runs optimization to reverse the flow of the model. At every step, we encode the canvas, compare it to a reference concept (for instance, how much does this image embody “dogness”?) and then optimize the image to embody more of that concept. Once again, we are modifying reality per se. The resulting images are quite reminiscent of psychedelic visuals, possibly hinting at what happens when top-down processing overpowers bottom-up

This is why outputs of models are often referred to as hallucinations. Not just when they hallucinate the wrong answer, but more generally the entirety of their outputs begin from nothing or partial context, and then are built by by a cycle of reading the current state and making a guess that is written back onto the canvas.

Hallucination is a feature, not a bug. It is the whole means by which we are able to generate new content. It is only problematic when the continued sentence by the language model tells me Paris is the capital of Antarctica. However, when the model correctly tells me Paris is the capital of France, it has still hallucinated one of many possible text continuations, just one that happens to align with our reality. Therefore, the holy grail lies in achieving controlled hallucinations.

However, when we feed an image to a visual language model for it to answer questions on, the visual signal itself is static. While the text thereafter is hallucinated, the image itself is processed entirely bottom-up.

How could we achieve top-down/generative recognition in AI? What are the advantages?

A few thoughts come to mind as to how we may develop recognition models with top-down processing.

This is a prime example of where we’d want to be cautious. Upscaling a heavily compressed image of someone reveals the model’s priors/biases, biases that do not align with our own. With a popular image of someone famous like Obama has a high level of expectation/bias given the iconicity, plus perhaps the fact that the features we most depend on for recognition may remain more intact despite the destruction, the distorted signal is enough to trigger proper identification for us. However, for the model, imaginably trained on a dataset of disproportionately white people, is biased in a different direction.

Now the funny thing is with guesses from partial, compressed, or distorted information like this is that there is no right answer. For all we know, both we and the model are wrong and the distorted picture of Obama is actually a close Obama look-alike. When going from low information to higher information content, we inevitably have to fill in and hallucinate the gaps. It is a one-to-many problem in that there are many possible high information data that when compressed could give the exact observed downsampled version. Every person and AI model has different inductive biases as to what they will assume, and its quite possible someone who has never seen Obama could make a guess of something like the example on the right.

Confidence

Where this becomes appalling is the confidence the model appears to show. When people don’t immediately recognize a downsampled image, we might admit, we don’t know who it is. Instead of jumping to a conclusion, an exact sample, we may sit in the superposition of all the possibilities it could be without actually committing to any answer. The model on the other hand, whether a stochastic or deterministic upsampler, is asked to deliver a sample instead of a superposition. I would equate this to being given a downsampled image of a person and someone forcing you to make a high-fidelity sketch of what this person could look like. You’d have to rely on your biases to fill in the gaps.

A lot of models are actually capable of exposing their uncertainty, its just not often revealed. For instance, the sharpness of a VAE distribution or a categorical distribution for discrete reveals something of the model’s confidence in its answer.

Nonetheless, it feels important to use confidence/uncertainty measurements as a guardrail against incorrect hallucinations. This is some of the philosophy behind projects like Entropix which use the entropy of model predictions to choose what words a language model puts out.

A Common Language to communicate by

One neat thing about sharing similar biases and priors is that it enables communication of data that would be stored in a large number of bytes but as very compressed versions. Despite the highly compressed Obama, we, the humans, got the message across, and you could probably compress the image quite a bit further. A similar case for the photo up by the title of this post.

We can speak in highly compressed codes because we know what the code represents. This is also why we can convey rich ideas in the language we speak, but are clueless to languages we are unfamiliar with. More subtly, when it comes to slang or words with multiple meanings, people of your in-group may immediately get what you mean, but those with different biases/expectations might understand your speech to mean something else. This is one place humans may still be able to beat machines on Captcha’s. We share biases and experiences that we are yet to crack with AI, and thus gravitate towards a common “answer” or that is not an obvious choice for machines.

This was especially apparent for problems in the ARC-challenge set. As a primer, given the sequence 1,2,1,2,1,2,_,_. what would you imagine comes next? Probably 1,2. Though this is not a “correct” answer. The masked portion of the sequence could be anything, maybe it starts going 1,2,3,4,1,2,3,4 and looping that until 1,2,3,4,5,6. A new pattern. But we are all commonly biased to recognize the original a reasonable solution, so much so that it feels obvious. The ARC challenge was similarly interesting in that AI got some problems right that appear very difficult to humans, and some wrong that appear obvious to humans. We differ in how the solutions we gravitate towards in ambiguity. Humans are right in similar ways and wrong in similar ways to other humans, emphasizing the common mental models we share.

I would go as far to say that these specific biases, what’s hallucinated in every scenario, embodies not only the remaining gap between AI and humans, but also what differentiates humans between each other and makes each one unique.

Multimodal LLMs hallucinating in the image space.

Note, this is more fun speculating than anything, I don’t have a ton of intuition to go off to imagine these approaches would result in aligning with human biases and solving our problem.

Normally, multimodal LLMs receive clean images, in contrast to the diagrams of human vision which has a narrow focal point. So perhaps the opportunity to make use of hallucinated vision isn’t really present in the first place given the image is already fully intact.

So, it would seem we’d want to somehow perturb the image + give the model a chance to restore the perturbation before generating text. The kind of noising we choose biases the model into relying on certain features, i.e. gaussian noise wipes out higher frequency details forcing dependence on lower frequencies. In the limit, one could craft a perturbation strategy that would aim to mimic human vision by giving it the same flaws (but also need to give it the same pros somehow as well? Like boosting outlines or contrast?)

Though, performing something like masking and prediction or noising and denoising with a separate model just to feed it through to the VLM afterwards just feels like a data preprocessing, or at the very least, somehow this would need to be jointly trained and part of the same objective.

Somehow, what I’d like to do is create a sequence of refinements preceding the LLM’s answer. Something like this

At train time, only the noisy observation would be shown. In place of noising, one could also remove patches of the image, or apply blur around a focal point.

I’d like to somehow also encapsulate a looped reasoning process, where we alternate between a text hypothesis and an image hypothesis. To do this, I think we could rearrange to move text to the middle. This can then be looped by truncating context, noising the new reconstruction, and then let that serve as the noisy observation.

What I would imagine doing here is looping this process until we reach a steady state where neither text nor image end up changing much, this would be perhaps a sign that generative processes and recognition processes are in equilibrium. One might notice that this somewhat resembles a VAE where everything is in context. A noisy observation is encoded into text stochastically and then stochastically decoded.

One issue here though is I could imagine potentially veering off away continually rather than converging. To ensure we have a way of anchoring ourselves to the original data, we could potentially show two noisy observations:

• one which is the original noised data

• one which is the most recent noised sample

Multimodal LLMs hallucinating in the text space.

The second approach I would imagine is... already happening! When you give a VLM an image and have it reason in text space before giving an answer, I would argue that generated text serves as hallucinating on top of the original visual observation.

As more text is generated, later text then conditions on both the original image, but also previous text it has seen.

If we imagine prompt injecting “This may look like a cat, but don’t be fooled, its a human in a costume, consider this when delivering your answer” this would inevitably affect our final output. Or before receiving the image, imagine prefixing the text, “You are an artist looking for something to paint” or “You are a historian looking through photos of the 2000s”. This also may skew the intrerpretation.

I would call this “External Hallucinations,” where text is injected into the reasoning chain influencing how the model interprets the image.

If we imagine having a multi-turn conversation with an LLM, giving it an ambiguous picture like the Downsampled Obama, it might volley ideas at us as to what the picture contains and be met with our feedback. After enough back and forth, its answer can be influenced by the picture plus the whole script of text containing both its own generated text plus ours.

A criticism with this approach is the actual image features, when encoded, are entirely deterministic and invariant to anything else the model writes after.

As stated previously, I don’t know if this is exactly a desirable trait, nor how much these methods will give us what we’re looking for. After all, the first method I could imagine possibly harming performance. Though it would be really interesting to know what it would take to get a model to recognize the picture at the top of this without seeing something like that in a dataset.

I’ve seen a lot of convoluted tutorials on attention but nothing really made it click for me more as understanding as mixing a projected version of your tokens across a sequence weighted by their “softmaxed” unnormalized correlations in a different space.

Let’s strip that back a bit and assess each part

• Projected version of your tokens

  • We’re not mixing the original copy of the tokens, but the Value projected versions.

• Unnormalized

  • correlation is measured by dot product, with is unbounded as opposed to being in between [-1,1]

• Correlations

  • the NxN matrix of the dot products between every vector combination. These are the mixing weights, for new tokens to be created as weighted sums.

• In a different space

  • tokens are first projected into another space by Query/Key projections when calculating the correlations

• Softmaxed

  • As opposed to using the values of the unnormalized correlations directly as is, we apply the softmax kernel which normalizes and “sparsifies” the weights.

Let’s try mixing your tokens across a sequence weighted by their unnormalized correlations. Sans the “different space”, “projected version of your tokens”, “softmaxed” parts. We’ll then build up to attention from there.

If you’d like to follow along, all experiments are performed in this repo.

1) Boneless Attention. Just correlation

Before drawing out this minimal attention with matrices and vectors, let’s look at it in a symbolic way.

Each one of these arrows/vectors is a token. The vector is in 2d space and can be described with an x coordinate and a y coordinate. The full set of vectors is our sequence.

Now, we’d like to create a new set of arrows aka vectors aka tokens, by mixing our existing ones together, but the ratios we mix them at should correspond to how similar or related they are. How can we measure how correlated two tokens are? By how closely they point in the same direction.

Here, we’ve plotted the similarities of the first token with the others. The opacity of the red arrows represents how similar the first token is with the one it is pointing to. Note, we are also having it point to itself.

Now with our similarities in hand, we can mix the tokens to obtain a new token corresponding to the first in this sequence.

We can illustrate that a given token should contribute less to the output by making it smaller. When we add all the vectors up tip to tail we’ll see it has less of an effect.

We then repeat the process for all the other vectors in our sequence.

Now we’re ready for a bit more mathy notation.

First, we need our correlation matrix.

Next we will “mix” our tokens, as a weighted sum with weights determined by the values in the correlation matrix.

Because we have 3 vectors, each new vector will be a mix of those 3. Looking at the first one (on the left cell of the blue row), we have:

Now doing them all simultaneously, as a matrix multiplication.

Notice that the vectors didn’t change a ton, aside from changing length a bit and slight shifts to direction, a product of the fact that they highly interacted with themselves but not as much their peers.

If we repeat this process over and over again, vectors that are less than 90 degrees apart from each other will increasingly move closer to each other.

This stripped back form of attention, but also attention somewhat in general, I like to describe as a “consensus” algorithm. It encourages the vectors to all agree on a direction, which is a smoothed, mixed result of all the initial vectors. This is a parameterless (boneless) attention. There are no learned weights.

However, this explosion of length makes for not such an exciting operation.

Next up, enter softmax.

2) Boneless Attention. Softmaxed correlation,

Same as before, except now we’ll apply softmax normalization to the correlation matrix. This makes the weights sum to one, avoids negative values, and also has the capability of “sparsifying” the operation, often driving many values close to 0.

Now, despite having no actual parameters, I’d argue this confers some practical value.

One example is where I used parameterless attention as a clustering algorithm for colors in an image.

Another, is that you can actually train transformer networks like language models using this kind of layer. The performance is not nearly as satisfying, but its also very far from terrible, raw correlation alone is at least doing something. (this is 1000 steps on GPT2 just as a proof of concept, so take with a massive grain of salt)

Now, let’s add in the Value projected version of our tokens.

3) Attention. Softmaxed correlation, weighted sum over Value projected versions

Here, as opposed to using the same versions of our original vectors for both determining correlation weights, and what tokens get mixed together, we’ll use a separate set for the vectors to be mixed.

This can give us a bit more of an uplift. Interestingly, loss plummets faster initially, perhaps because correlation in native space as opposed to one created by random projections may be a stronger starting place off the bat. But as Query and Key weights are given time to adjust, baseline quickly surpasses this modded version.

4) Bonafide Attention. Correlation can be taken in another space.

Here now, we introduce the Query and Key learnable projection matrices. When we measure our correlations between tokens now, we’ll do them between two different versions of the original set.

Why would we want to do this?

One of our original problems with attention is that it is guided towards consensus, but not too much beyond that. We also found that the self-interaction was also always very high, sort of loosely explaining part of what the outcome will be before we even proceed.

By projecting to another space, we can more finely decide which tokens should correlate with each other, increasing the expressive power of attention.

This operation is WAY more flexible and expressive in higher dimensions, but we can imagine this simple 2d case where we have 4 vectors. We could apply a stretching transformation as such, increasing the similarity between nearby vectors and further decreasing it between ones farther away, letting us modify correlation weights and thus the output of attention.

Now for fun, let’s see what happens if we double the initial values of the correlations prior to feeding them to softmax. This should help give a feel of the advantages of softmax and its sparsifying properties. For some more interesting tidbits see here.

This then becomes the weights for the value tokens (green in previous diagrams) to be mixed by.

At this point, the tutorial is complete, this is attention! The remainder of this post is further experiments on attention leveraging the fact that it is somewhat similar to the idea of calculating raw correlations.

5) Bonafide Attention. Initialize Query/Key projections to identity function

Seeing the initial fast loss drop of the previous method, it seemed worthwhile to try initializing the Query and Key projection matrices to the identity function, BUT unlike pure “just correlation” attention, they are free to adapt from there. Perhaps this is a good place to start and adapt from?

This might be our strongest contender so far, though still baseline surpasses it.

I Imagine a counter argument here is that this kind of thing simply doesn’t optimize well, and we do our random initialization for a reason.

There’s still more we can try though. What is a common way we can include the identity function? Residuals!

6) Bonafide Attention. Make Query/Key projections residual layers

Here now we change the query and key calculation from

Q = wQ(x) K = wK(x)

to

Q = wQ(x) + x K = wK(x) + x

This is motivated by correlation making for a useful starting bias, but also wanting to give adequate freedom to the optimization process.

Here the two were now run for about 4k steps. The two are approximately neck and neck with residual attention seemingly having a slight edge.

A Discussion on Low Rank Attention Weights

A point I’ve seen a few papers contend on is the rank we observe of the weights used in attention layers. This is approximated by taking the singular value decomposition and seeing how % of singular values are “properly large”.

Low rank layers are generally seen as a bad thing. It can suggest dimensional collapse and an underutilization of the network’s capacity. But is this always true? When we begin training, sampling our weights from normal distribution, our matrices begin as 100% full rank, so it can really only go down from there regardless.

An instance here is the the knowledge that CLIP embeddings are low rank, a few axes explain the majority of the variation in the output, making them feasibly much more compressible and lower dimensional than they reveal. Here we have the singular values extracted from the covariance matrix of 500,000 OpenCLIP VITG embeddings

It very rapidly approaches 0

We can also see how this affects the output by decoding the embeddings back into images, after pruning out smaller singular vectors. This is an example from the DALLE2 paper.

Starting from the top 320 dimensions, we can prune out the last 120 with not a majorly substantial change to the output. Though once we get down to 40, 80 range we can see these larger dimensions are much more vital.

The Modality Gap paper looks at CLIP in depth and the factors around contrastive learning that induce this low-rankness

This paper also then finds a way to address it, recovering full use of dimensions and slightly increasing performance.

So at least for embedding representations, low rankness seems undesirable but evidently far from catastrophic.

The paper, Weight decay induces low-rank attention layers, comes off as a scary title. However, the results are a bit surprising.

They measure low-rankness by the observed 95% saturation point. We can see that higher values of weight decay, indicated by lambda, very much reduces the rank of attention layers. I’m tempted to say the the charts on the right are particularly interesting, that the pseudo rank drops and then increases over the course of training, however this measurement doesn’t imply that dimensions have fully collapsed and then recovered, it could just as much be our largest dimensions decreasing as small ones increasing. Nevertheless it’d be interesting to see if that chart continued a bit more to the right.

Despite finding this massive drop in rank, performance does not plummet in the same manner. In fact, performance seems to improve a bit in some cases, and this effect extends to MLP layers as well. The benefit could be explained by low-rankness, or it might be some other factor around weight decay shrinking the magnitude of weights offering regularization. It could also be the case that “pseduo-rank” might not be telling us the whole story. Regardless, it doesn’t correlate to the expected loss of performance we would imagine with decreased rank.

So perhaps, this is not something we should aim to avoid but something we can lean into? Maybe attention layers, as they are, are over parameterized by dense matrices, but possibly this suggests this is a good place we can use structured sparse matrices like LoRA.

And this is kind of exactly what Deepseek’s Multi-Latent Attention does, both reducing computation a bit and the size of KV caches.

This was a similar motivation for the experiments on correlation done in this post. If correlation can already do alright, do we need such a heavily parameterized operation?

Calibrated Attention

Calibrated Attention NanoGPT

Attention is the magic ingredient of modern neural networks. It is the core of what has launched performant language models into the spotlight starting with GPT, and since then, it has extended its hands across all modalities.

There are a number of desirable properties that make attention a first-class building block. Namely:

• It handles variable sequence lengths with ease

• It allows for a global receptive field without needing to scale parameters

  • as opposed to convolutional neural networks

• It enables dynamic routing of information between tokens

• It is parallelizable and incredibly GPU-friendly

The core principle of projecting tokens into another space, calculating their correlations, and using those as weights to define mixing strategies of the values is solid.

But here I want to very boldly claim that our specific choice of the well-known attention equation is a fluke. In particular, softmax appears to come out of nowhere.

Why Softmax?

Numerous papers have developed alternatives that did not end up making it to the mainstream for lack of competitive performance. So why dare to make the 1001th attempt to contend with this? Is it possible we got it right the first time?

To be fair, most attention alternatives aim to circumvent the O(n^2) scaling of computation in sequence length by attempting to linearize the computation or use cheaper, more flexible activation functions. Attempts to skirt around the cost may naturally lend themselves to drops in capability. Perhaps the O(n^2) is inevitable.

However, let’s stay with the core circuit of attention. The part I’d like to question is why specifically the exponential function, normalized across the sequence.

It’s a natural choice. Softmax appears all around neural networks. It is a sensible means of turning real-valued vectors into a row of probabilities. It is also nicely differentiable and has the ability to non-linearly “silence” and filter out noise to reveal signal.

Attention, I would argue, makes good use of the latter traits. Here’s an example of the suppression it can provide.

Let’s imagine here that we have taken one vector and dot’ed it with 10 others, loosely measuring their correlation scaled by norms. This is a plot of their interactions.

If we were to take a weighted sum of the 10 vectors with these coefficients, we could end up with a vector whose norm is quite large, where we generally like to have a consistent steady range of values for our neural network to run smoothly and not explode. So we can start by subtracting by the mean and sum normalizing it.

This is okay, but note that our final output will be a pretty mixed blend of all the tokens. It’ll be 30% token 6, approximately 15% token 7 and 15% token 8 and so on.

But let’s say we wanted to pull out the token with the largest interaction, or maybe a few, and not have noise from the other tokens mixed in. This is a common desire in database lookups (an established analogy for the role attention and dot products play in neural networks) where we might take be interested in the top-k vectors.

We’ll want some means of setting the other positions to zero, but just doing max(x) or manually setting other positions to zero is not differentiable. Enter softmax.

Now if we take a weighted sum over the tokens, we almost recover the function of max(x). The output will be approximately ~99% token 6. This function has a useful winner-take-all nature to it via the nonlinear exponentiation followed by normalizing across the sequence. The growth of mass in one value means a loss in others, even being able to drive its comrades close to zero.

I’ve made a case up until now for why we really like softmax. It’s noise suppression is useful. Though the other aspect of softmax is its root in forming probability distributions. For that, there isn’t exactly a reason we need to think of attention mixing weights as probabilities.

There are plenty of other kernels that can yield a similar effect. For example, x^3 actually gives a pretty similar result for this case

This is the argument of this work, which reveals a number of polynomial kernels that can also perform well on toy tasks.

Additionally, do we need mixing coefficients to be exclusively positive? Why?

This was an argument I made here, where I argued that negative correlations may be just as much informative as positive ones, and it is perhaps the small correlations instead that should be suppressed. This later became a paper under a similar motivation.

Lastly, another work explored breaking free from typical normalization, allowing attention to take weighted sums outside of its convex cage.

So the questions I see are:

1. Why e^x? Especially if we don’t need to use the values as probabilities?

2. Do we need sum-to-one normalization?

3. Why only positive values?

None of these, however, feel like bugs.

We may find further merit in searching the space of alternatives that deviate from these conventions, but after all, as I said, softmax was a fair choice for a first iteration due to its common usage and it was sufficient to bring us where we are today.

However, there is one aspect around how normalization and scaling is treated that I do feel is a bit of a bug with attention. Particularly in cases where our sequence length can vary, like with causal language models.

Normalization Across Variable Length Sequences

Loosely summed up in this tweet here.

The one part I find particularly fragile about attention is the dynamics around normalization, and how it can vary depending on the length of the sequence. Firstly, let’s see a toy example of this in action. We’ll sample logits randomly from the Gaussian distribution across different lengths and apply softmax to them. Observe what happens to the size of the largest value, entropy, and overall sharpness of the distribution.

At 2 elements (Note: if we only had one element, 100% of the mass would go to it always, it would be a weighted sum over 1 element making attention pointless)

At 5 elements

At 25 elements

At 1000 elements

The distribution becomes progressively more flat, entropy increases, and the largest value becomes much smaller.

This is not just a reality of Gaussian variables, but is also observed with attention maps from our models for a generated piece of text, namely GPT2-large here.

At length 5

At length 50

At length 250

Ignore the first token for now, that’s something we’ll talk about later. With growing sequence length so does the requirement to produce outlier-ly large values to overwhelm all the others and achieve sharpness. We end up with many small spikes rather than few large ones. Now entropy naturally increases with the number of choices, though again, is it necessary to apply the probability lens to attention?

It took reading a few works before what’s happening really cemented in my mind, and now it’s impossible to look away from.

Achieving Null Attention

The first was Softmax-plus-one, which argues attention should be able to do “nothing.” I don’t consider this a huge issue, and not the aspect I think of as the aforementioned bug, but a solution here would feel like a “nice to have.”

This blog post criticizes the dependence the denominator has on the numerator. It could be the case that the a token would like to choose inaction, or have weighting coefficients that sum to less than 1 to reduce the magnitude of effect of a given attention layer.

Recall softmax looks like this

In the current formulation, this isn’t particularly possible*. We can try to reduce our token-token interactions as much as possible, even making them exclusively negative, and attention would still have an effect. In fact, one could imagine an attempt at this kind of adaptation could backfire. If all correlations drop evenly, we’d approach a uniform distribution of attention as opposed to isolating a few candidates, which is exactly what brought us to using softmax in the first place!

The “fix” is to add a bias to the denominator, slightly decoupling the denominator’s dependence. It would now be possible to output very small values in the numerator but have the denominator remain as a constant value.

If we imagine that

Then

Becomes approximately

Allowing for the capability to assign 0 weights everywhere, entirely nullifying attention. Though this can also appear more partially where exp(x) airs around small values below one allowing us to say, have our total weights sum to arbitrarily something like 0.2.

I have yet to see this implemented anywhere, and in my opinion the missing piece to practicality is that the denominator bias should be learned rather than fixed to 1. 1 is an arbitrary choice and may ask attention weights to output around a certain range of values that may be unnatural for it. (though this too, as we’ll see might not even be sufficient due to how the denominator’s size also depends on the exact length of our sequence as well, we may prefer a learned bias that is a function of the sequence size)

But why would attention want to do nothing? I paid for the layer; am I wasting my FLOPs?

Inaction is actually a good bias to have in neural networks in general. It is often a good heuristic to firstly assume the identity function, no effect, and then increase the effect from there. This is partially the principle of why we use residuals everywhere nowadays; learning the identity function directly by gradient descent is actually quite difficult. We can instead hardcode the identity function and ask that the networks learns to predict what should be added to it, the residual.

I left an * above on the remark that inaction is impossible for attention layers. Despite a lack of conventional means of choosing to do nothing, we’ll often find that nature (neural networks) finds a way anyway in a form that is equal parts clever and aggravating. What you’ll find is that tokens will often put a really large portion of their attending mass to a meaningless token like a padding, or a comma, or beginning of sequence marker. This paper aptly referred to this phenomena as “Attention Sinks.” A token is forced to attend to its bretheren up until it fulfills the entirety of its attention budget or mass. Though it can choose to just dump the majority of it onto a token that confers no effect.

I found similar effects very extensively in text2image models and a bit also in text encoding models.

This shows the attention of the image to tokens of the text prompt, where y axis is the amount of attention mass and the x axis is the token index. Here we are displaying, for every text token of the prompt, what is the maximum amount of attention it has ever received from the set of all image tokens. Across different attention heads

We can see there is a massive amount of attention on the very first token of the sequence which is the BOS token. Importantly, because CLIP here encodes the text prompt with causal masking, the BOS token is entirely invariant to the text prompt itself. It is the same regardless of what you put in and thus confers no information about the request. So the remaining hypothesis is that similar to Attention Sinks, we are attending to “dummy token” just to reduce the weights we place on the other informative tokens, reducing the magnitude of the output, which we would probably rather have a mechanism for scaling more directly!

Incredibly interestingly, we can also look at the norms of the projected values of all the text tokens. The norm of the BOS token approaches 0 meaning it gets quite close to being a vector of all zeros. If we place the majority of weight on this vector of zeros, our resulting output is pretty close to a vector of all zeros as well!

Though I cannot help but feel this is maladaptive in the ways this requires the value projection matrices to adapt to the BOS token vector, and that there are better ways to achieve this effect.

Sequence-Length Dependent Temperature.

The other paper that made me realize the importance this played was Training-free Diffusion Model Adaptation for Variable-Sized Text-to-Image Synthesis. And this where I begin to see a bit of a potential bug in attention.

Firstly, they reiterate our previous observations that the entropy/sharpness of attention depends on the sequence length.

However, if we train on only one fixed sequence length, a common practice for image models at the time, we might expect the model to adapt the norms of the Queries and Keys to account for this. But when we go to test out other resolutions, we’re now extrapolating to a domain we are ill-adapted for. Their proposal is to train on one resolution as normal, but then at inference time, rescale the temperature scaling to attempt to match the expected entropy we would have at our learned, trained resolution.

However, the need for a relative anchor makes it a bit of a band-aid approach and its common to train across multiple resolutions simultaneously. We can’t match our scaling to that of a trained length if this is a moving target. We’d really like to have this built into the attention equation itself at train time.

On the topic of handling multiple sequence lengths simultaneously, language models, which train with causal masking, zero out the logits of tokens to the right (future) for a given token.

So technically, if we’re training at a sequence length of say, 4096, attention is effectively batching queries that attend to a sequence of length 1, a sequence of length 2 ... all the way to a sequence of 4096 simultaneously.

As stated before, our query and key projection matrices can adapt themselves to output queries and keys that have norms allowing to achieve desired levels of entropy.

But here’s the problem, the same matrices are being used for all tokens. Some sequence lengths may prefer larger norms and some smaller, but we can’t adapt in both of these directions at once over these single shared weights! This is not to say we need a set of weights for each possible length, but there may be something we can do.

Hypothesis: This may be correlated to issues of entropy collapse and instability in attention layers seen in LLM training. The sequence of 4096 may request a gradient direction leading to lower entropy via higher Q and K norms while the sequence of 5 may request higher entropy. Conflicting update directions may guide training towards a poorly adapted and fragile compromise.

A proposal:

1. The scaling factor of attention should be a function of sequence length.

   1. This could be learned or scaled by an equation, leaning towards the equation approach as it seems there is a good way to understand this.

2. There should be a bias in the denominator that is either a learned constant, or possibly also a function of sequence length.

   1. This I have less intuition on though I lean towards a learned approach as it may depend on the typical range of values the logits take on.

With respect to 1, Petar Veličković had a paper on this very matter, for how attention varies with sequence length and how we can recover sharpness through scaling.

In conclusion, while we have now detailed explanations for how attention performs after the fact, its inception was somewhat a fluke in terms of justifying some of the specifics of its design. A common and effective trend in ML, to try first and find out later.

Attention is obviously performant as is, but it appears there’s potentially some places to grow if we forego the probability framework softmax lends itself to and possibly resolving something I would even deem a bug.

@misc{smith2024attndiffusion,
  author = {Smith, Ethan},
  title = {Softmax Attention is a Fluke},
  url = {https://www.ethansmith2000.com/post/softmax-attention-is-a-fluke},
  year = {2025}
}

A short attempt at a small portion of the diffusion Family Tree

https://www.canva.com/design/DAGgnVB3x2s/b52Y3Kg-frWdRlPzI3_5pA/edit?utm_content=DAGgnVB3x2s&utm_campaign=designshare&utm_medium=link2&utm_source=sharebutton

Discrete diffusion is the lesser known (but growing in popularity) sibling to the familiar Gaussian diffusion models. Despite their differences, these two both fall under the family of models defined by:

• A random iterative destruction forward process. When reversed, it allows for generation.

• A model that either directly or indirectly learns the score of the distribution, the “direction” to move in to increase the probability of a sample and move to a less noisy state.

Continuous diffusion typically looks like this:

https://cvpr2022-tutorial-diffusion-models.github.io/

Discrete Diffusion you can think of looking something like this (masked diffusion case, and typically this is with latent tokens passed to a decoder instead of being done in pixel space)

Discrete diffusion seems to be the accepted name now for a large family of methods that involve probabilistic removal/replacement of local patches of a piece of data and using a model that predicts a categorical distribution to recover the original clean sample.

Gaussian diffusion involves globally interpolating a piece of data with Gaussian noise, such that the result is a mixture of noise and signal. At higher noise levels, noise dominates while signal is sparse, and vice versa for the lower noise cases. When we normalize both data and noise to std=1, we can measure how noisy the data becomes by its signal-to-noise ratio. We have a continuous state describing the mixture between signal and noise.

On the other hand, instead of a continuous signal, like pixels, discrete diffusion instead works with a categorical vocabulary. Let’s say, we have a vocabulary of 8192 unique tokens. We encode our data and assign patches of it to one of these tokens, by an approach like VQGAN.

Now it doesn’t make sense to add continuous noise to integer values that simply indicate the index of a vocabulary: nearby entries in the vocabulary aren’t necessarily related.

So now, we have two options for noising. From a single token, we can either assign it to a special value that never occurs in the data, like [MASK] (absorbing end-state diffusion), or we can randomly switch it to another entry of the vocabulary (stochastic end-state diffusion?).

What’s interesting here is that noise is no longer a continuous spectrum but rather a binary variable: the token is either its original state or the noised state (MASK/random value) it got swapped with.

This introduces a challenge in adapting from continuous diffusion. We want to have many different noise levels to perturb and predict from. Instead, we can think of our noise level as representing a probability that an individual token gets swapped with its noise copy. So at noise level 0.5, we’d expect half the tokens to be swapped.

Notice now that the noise level is no longer consistent everywhere. With our Gaussian diffusion, the same amount of noise is applied everywhere around the data (granted, we may sample values closer to 0, so this in a way constitutes less noise). Here, however, a portion of the image may remain entirely intact while other tokens have been swapped.

An ongoing challenge of discrete diffusion is actually identifying what kind of SNR we are at. Let’s imagine, for instance, in the top right corner of every image is a special token that indicates not information about the specific patch of the image but instead classifies the image globally, like “cat, dog, horse, car...” etc. Noising this token may have a different effect from removing other tokens, so it’s hard to say what canonical noise level we are actually at. Another example of this is with diffusion language models. Removing tokens like “the” or “a” does not constitute a huge loss of information: they are easily predictable from the grammar of the sentence, and the message of the text is well preserved without them. Other words of a sentence may be much more important, and their loss would confer a more substantial loss of information. This is also a reason the MASK option is preferable to the random replacement case. If we happen to replace a word with a synonym, this may not confer too much loss of information, but a more out-of-place word could result in more significant destruction.

Once upon a time I thought I had come up with discrete diffusion when I was seeking a way we could reduce the neural network modeling load by augmenting it with wave function collapse. In particular, I was interested in ways of building 3D models from learned super-voxel tokens and having some kind of ruleset for how they should be built.

I called the approach Wave Function Collapse BERT for Images.

Wave function collapse is used for procedural generation, often for something like randomly generated dungeons in games. There is a “hard” rule set that says as one goes about placing different kinds of tiles, only certain tiles are allowed to neighbor them, thus reducing the uncertainty of what remains.

For instance, the image on the left shows our fresh grid. Every tile currently can assume any state because there are currently no constraints. However, once we place down, say, a tile with a barrel on it, we know that just below it needs to be a rock. On the right, we can see that the options for what we can choose have been reduced.

I was quite fascinated by the results of autoregressive image generation at the time, but it takes a criminally long time to place all of the tokens of an image (which may easily be in the 1000s) one by one. It was also strange to me the ordering we placed tokens in. Why do we start at the top left corner and place tokens in raster order? This could allow us to sample multiple tokens simultaneously, especially as the context becomes more clear, more and more tokens become “free”, similar to how we have fewer options to choose from as more of the dungeon tiles are filled out.

Pretend in these diagrams instead of sampling from words, we are sampling visual tokens.

Discrete Diffusion is kind of like wave function collapse where the rules are “soft” (meaning that which tokens are allowed to coexist with each other is based on probabilities of what occurs rather than absolute rules) and are learned by the model. Because of mechanisms like attention as well, we can have learned rules that have global awareness of the board as well.

Sudoku is an example of a problem that can be solved via wave function collapse, which made it seem like a fun thing to try with discrete diffusion. I observed that this was able to learn to predict valid sudoku boards starting from scratch with as few as 2000 training steps.

https://github.com/ethansmith2000/DiffusionSudoku

It’s a bit hard to see in the diagram but in addition to being convolved with a gaussian, these points are also drifting towards zero.

There’s two perspectives here actually. There’s setting a point to

xt = x0 * alpha + noise * sigma

where sigma and alpha are both numbers between 0 and 1

and then there’s

xt = x0 + noise * sigma

but sigma goes towards infinity at the end of the diffusion schedule.

In both cases we can achieve a desired signal to noise ratio but one case involves reducing the image signal while the other keeps it constant and continues raising the noise’s variance to overwhelm the signal entirely. I believe these are the Variance Preserving and Variance Exploding perspectives respectively.

We go from a very high spread distribution to something narrower and narrower until it is a mixture of diracs each reflecting a single sample.

Interestingly, this process is opposite to entropy flow in nature and the arrow of time. Typically entropy, or disorder, increases. When you put a drop of coloring in water, it will slowly spread out to become a homogenous mixture. The way this moves forward in time is a stochastic process, but the end result distribution is well understood, even if we don’t know the exact end state it will end up in. In all cases the end result appears similar.

Now imagine if we could reverse this process. Could you guess from a totally homogenous mixture where in the glass the coloring was dropped? No. But we can guess at this. Generative diffusion is estimating that guess based on where it could have happened with greatest proability.

I put together this repo, Boneless Flow, to explore it some more. Instead of training a model, with weights, to obtain estimates of the flow towards the manifold of clean data, we can compute the ground truth flow analytically if we have the whole dataset in memory.

A problem with this is that the ground truth score actually won’t allow for generating new samples

https://arxiv.org/abs/2412.20292

Tweedie’s formula is an equation that sits at the foundation of diffusion. It describes the visual above of moving towards regions of higher probability.

I liked this post on it. Though I wanted to give some additional intuition.

Here X is a clean data point, e is noise, and we get out Y when adding that noise to X.

f (denoiser) is a function that takes in Y, the noisy data point, and attempts to estimate X, what it was before being noised.

f_hat is a specific instantiation of f namely, the MMSE denoiser.

It says that we can denoise Y by starting from it and moving in a direction that increases probability, similar to our original diagram, scaled by sigma, a known level of noise

We can consider an example of this without needing to consider diffusion: Linear regression

In a linear regression problem we have a set of points following some kind of trend, we attempt to estimate this trend with a model, a straight line represented by Y_hat = mX + b.

Y_hat is the predicted value from our model.

Y is the actual data points, which we can think of as clean data points from our model but with added noise. This is often a reasonable assumption due to other unaccounted for variables, which we can represent as noise.

Therefore we can think of the Y‘s represented as

Y = Y_hat + e

or

Y = mX + b + e

This noise, assuming gaussian distribution, creates a fuzzy field/distribution around our model, which we can visualize as either of these two. Where height of the distribution represents probability density, or the opaqueness

Now like before, we can ascend the probability curve, thus denoising the real data points with respect to our model.

I think this illustrates here a profound relationship between noise and blur, for which blurring can be seen as denoising.

Noise is individual random samples from a distribution.

Blurring involves taking the average of many samples, or if we know the distribution, typically pointing out its mean value.

We combat noise by blurring, like this:

We don’t know what the image looks like without noise, but we can take advantage of the contiguity of images to assume that nearby points should have related values, thus smoothing out the randomness. (note this is spatial blurring, which is a bit different from the kind of blurring regression performs, but both leverage the aggregate of noisy samples to give us a denoised estimate)

Blurring gives us the best unbiased estimator of the data without noise. It yields points that have minimal distance to all the observed noisy samples. The unbiased choice is also the most conservative, safe choice: it makes minimal guesses.

We could also average or blur across other axes of data points. For example, instead of blurring spatially individual pixels, we can treat every picture of a face as a sample and average those together.

Note, this doesn’t actually give us a real sample of a person, but its informative of common traits in appearance from a certain cluster. This is reflective of typical samples vs the mean sample.

https://sweet-hall-e72.notion.site/Classifier-Free-Guidance-to-Approximate-RL-9f78c02801c6434da61f37c8d843c5bf

https://sweet-hall-e72.notion.site/Do-Diffusion-Transformers-Deserve-The-Hype-9b9ca7bead374b47aac96558714c203b

https://sweet-hall-e72.notion.site/Automated-LoRA-Discovery-and-Teaching-Neural-Networks-to-make-Neural-Networks-22aa3b5ad66e4bc985ff2c93896538d2?pvs=74

https://sweet-hall-e72.notion.site/Diffusion-and-Autoregressive-Models-for-Learning-to-Solve-Mazes-c3bc4bcdfa304ecd9531ee5445a4da66

https://sweet-hall-e72.notion.site/Traversing-through-CLIP-Space-PCA-and-Latent-Directions-b898932e13684d58957405b4a2747a79

https://sweet-hall-e72.notion.site/Learning-Space-Filling-Curves-with-Autoencoders-e39e41ce75894c3a8fecfee0f3bbfb23?pvs=4

https://sweet-hall-e72.notion.site/Mimicking-Diffusion-Models-by-Sequencing-Frequency-Coefficients-8e5a60e876d640c390369627d55330b1

https://sweet-hall-e72.notion.site/ContrastiveDPO-for-Diffusion-Generalizing-DPO-to-multiple-items-PART1-226b3746aa4d4ff9995d1e26b38a9674

Image from:

https://twitter.com/toshi2fly/status/911306344376012800

Post: https://sweet-hall-e72.notion.site/Dipole-Attention-Opposites-May-Be-Deep-Connections-6592f109c52c4caf97ffb5a3cd1fb45e?pvs=4