This is my submission to BlueDot's Technical AI Safety Puzzle #1, for which I received an Honorable Mention.

Congratulations to Gustavo Korzune Gurgel, Patryk Perduta (his amazing write-up), Sam Spilllard, Karine Levonyan, and Michael Zlatin for their recognition in the puzzle.

My article below focuses on my answer to Task 3: training a small MLP to encode country feature through a chosen nonlinear manifold in three reserved channels.

My Task 1 and 2 write-up is available on my homepage, and the interactive/more intuitive version of this article.

I welcome discussion, feedback, and collaborations that could extend this idea.

You can check out the code for this article in my GitHub repository.

The puzzle and the question

BlueDot's Technical AI Safety Puzzle #1 provides a trained five-layer MLP for multi-label classification over eight binary features, using mean-pooled sentence-transformer representations. The puzzle identifies nonlinear behavior at the output of the third ReLU, denoted as , and asks participants to:

• find the nonlinear feature ;
• explain the geometry used at to represent ; and
• train a new model with a more interesting representation.

This post addresses the third task. I train a new five-layer MLP and constrain country to use a chosen three-dimensional manifold while testing whether the classifier relies on that code.

The guiding question is:

Can we choose a representation geometry first, then train the model so that the feature follows that geometry while still solving the original task?

This is inspired by work on counting manifolds in language models [1] and manifold steering [2]. I use the opposite direction: instead of discovering a manifold after training, I choose a small manifold code first and ask the model to use it.

The manifold must not be decorative. I therefore ask both:

Did the hidden activations form the desired shape?

But also:

Did the model actually use that shape to solve the task?

My design has five requirements:

1. Specify the geometry before training.
2. Make the learned geometry code match it.
3. Preserve the original eight-label task.
4. Weaken easy shortcuts, especially linear and complement-only access.
5. Show that interventions on the geometry change the country prediction.

1. Define the Target Manifold before Training

The first stage has no model. I only define what the model will later be asked to learn.

1.1 The Distributional Problem

For each example, the dataset gives the binary label . It says whether country is present, but does not say where the example belongs on a desired manifold.

For a helix, supplies no position , radius , or tube angle . If those coordinates were observed, I could use point-wise supervision to train my model.

Instead, I use a distributional target:

Here, and are the learned geometry codes for positive and negative target examples. The model is free to choose individual placements; the aggregate positive and negative code distributions must match and . This follows the aggregate-distribution perspective of Wasserstein Auto-Encoders [3].

I choose two target geometries: a sphere shell and a helix tube, shown in Figure 1. In both, positives occupy an inner radius and negatives an outer radius while both classes span the same scaffold. The label is therefore a distance-to-structure decision rather than a global direction. If positives occupied only a sphere's north pole and negatives only its south pole, the representation would be essentially linear.

1.2 Sphere Shell

Let be a direction sampled uniformly from the unit sphere , the set of points in 3-dimensional space with Euclidean norm 1.

Let be a radius. The raw sphere point is

The positive class uses a smaller radius than the negative class:

with and . Here, is the target radius for positive examples and is the target radius for negative examples. The approximation means the sampled radius is concentrated near that value, not exactly fixed at that value. Thus country is encoded by an inner versus outer shell, not by direction.

1.3 Helix Tube

Let be a position along the helix and its vertical pitch (to control how fast the helix rises). The helix core curve is

Two local cross-section directions around that center line are

With cross-section angle and tube radius , a point in the tube is

The class-conditional distributions are

I use , , and . Both classes trace the same curved scaffold.

I tried three turns, but the tail layer did not reliably read them under the remaining constraints; one turn was learnable, usable, and testable.

1.4 Normalization

The sphere and helix have different natural scales, so I normalize sampled points to test shape rather than raw norm. Let be a raw point, with for the sphere and for the helix. With balanced classes,

For the sphere shell,

and for the helix tube,

Here is the radius-noise width. Its term is the uniform-interval variance rule [4]. The helix additionally contributes from its circular core and from its centered vertical coordinate.

Before training, I draw normalized target anchors and .

2. Train the Model

2.1 Base Classifier

Each text input is encoded by sentence-transformers/all-MiniLM-L6-v2, then mean-pooled into . A five-layer ReLU MLP maps it to eight logits and is trained with binary cross-entropy:

The ordinary classifier reaches mean AUC , mean accuracy , and country AUC . This establishes that later failures come from manifold constraints rather than basic architecture or data handling.

2.2 Add the Geometry Bottleneck and Activation-sign

2.2.1 Geometry Bottleneck

I reserve three hidden2 pre-activations for a geometry code , leaving 61 complement coordinates . Three channels are the minimal space for a point on either target 3D manifold.

Because hidden2 is post-ReLU, I add positive offset so a signed manifold can live in activation space:

Let  be the chosen geometry family, either the sphere shell or the helix tube. The learned coordinate head predicts three unconstrained values, and fixed map converts them into a normalized point on the chosen manifold:

For the sphere bottleneck,

where is the azimuth angle , is the vertical coordinate, and has maximum allowed radius .

For the helix bottleneck,

where , , and .

2.2.2 Experiment

At this stage, I optimize only . The purpose is not yet to force good class geometry. I ask only:

Can the classifier survive this architectural constraint?

Alongside mean and country AUC, I use four analyses:

1. Geometry probe. A fixed readout from alone measures whether its sphere radius or distance to the helix core is closer to the positive or negative target radius:

Its ROC AUC shows whether the reserved code contains the label in the intended geometric form; it does not prove the tail uses that code.

2. Coverage entropy. For the sphere, I bin the azimuths of learned points into eight bins. For the helix, I map each point to its nearest core position and bin it into 24 helix positions. High entropy means codes spread across the intended manifold rather than collapsing into one patch.
3. Causal target delta. I replace the geometry channels with positive and negative anchors, run the tail, and compute the mean country-logit difference. A large positive value means the tail listens to geometry; a near-zero value means the geometry may look correct but the tail ignores it.
4. Linear probe AUC. I train a linear probe on the full country activations. Lower is better: it means country is less exposed as an ordinary linear direction.

2.2.3 Results

Table 1. Model behavior after adding the geometry bottleneck. Higher is better unless marked ↓.

The model survives the bottleneck, but the near-zero causal deltas show that it does not use the intended geometry. Low geometry-probe AUC and coverage entropy also show that country does not yet follow the prescribed shape. The model is routing the label through other complement channels.

2.3 Add Unconditional Optimal-transport Coverage

Can the learned geometry codes cover the target geometry instead of collapsing to a small region?

The bottleneck constrains codes to the allowed family, but not to the whole geometry. Every example can lie on a valid helix tube while almost all examples occupy one segment.

To prevent this collapse, I add unconditional mixture optimal transport (MixOT), which spreads the unlabeled mini-batch across the overall target geometry:

where is the country-positive rate.

For mini-batch , let be its n learned 3D geometry codes and sample equally many anchors from .

Think of learned codes as students and target anchors as seats across the shape. OT matches each student to a seat while every seat must receive a match. Collapsed codes leave distant seats unmatched and are expensive; spread-out codes make the matching cheaper.

Therefore, for a learned code and a target anchor . The pairwise transport cost is

Entropic OT solves

subject to

Here is the transport plan. The first constraint assigns every learned point, and the second ensures every target anchor receives mass. This second constraint makes collapse expensive: the model cannot ignore most anchors and only match the easy ones. The entropy weight smooths the matching for efficient Sinkhorn optimization [5]

To test whether class-conditional geometry works on its own, I replace the previous MixOT term with:

2.3.1 Results

I add Radius MAE, the mean absolute error between learned and desired class radius (lower is better).

Table 2. Model behavior after adding unconditional mixture optimal transport.

MixOT makes coverage almost perfect, but geometry-probe AUC remains weak, radius MAE remains large, and causal delta remains nearly zero.

This is expected: it spreads the unlabeled batch but does not assign positives to and negatives to .

2.4 Add Class-conditional Geometry

Put positive examples in the positive part of the shape, and negative examples in the negative part.

MixOT spreads the batch but does not assign each class to its own part. I replace its shared seating chart with class-conditional matching and add local per-example signals.

2.4.1 Class-conditional OT (ClassOT)

For , let and sample anchors from . Then

where is the number of included label groups; I skip a class with fewer than two examples in the batch.

2.4.2 Radius Loss

ClassOT gives the right distributional shape, but a batch-level match can be weak local training signal. The fixed tail has to decode country from each geometry code. Radius loss supplies a simple cue: each example should reach its own class radius.

Let be distance from the origin for the sphere or distance to the nearest helix core point for the helix:

2.4.3 Geometry-score Loss

Geometry-score loss asks whether the geometry itself would classify a point correctly. Define

and the logit-like score

It is high near the positive radius and low near the negative radius, so I optimize

2.4.4 Geometry-score Loss versus Radius Loss

ClassOT decides where the two populations should lie. Radius loss asks each point to reach its assigned radius. Geometry-score loss instead asks whether the geometry itself would classify the point correctly: it turns distance to the positive and negative radii into a logit-like score, so a point closer to should look more country-like and a point closer to more negative-like.

For and , a positive at is good under both losses. A positive at is still pulled toward by radius loss, while geometry-score loss considers it acceptable because it remains closer to than .

2.4.5 Experiment

I remove the previous stage's mixture-OT term to test whether class-conditional geometry works on its own.

2.4.6 Results

Table 3. Model behavior after replacing unconditional mixture OT with class-conditional OT.

ClassOT writes country into the intended geometry but does not make the classifier read from it. Geometry-probe AUC reaches for the sphere and for the helix, but causal delta remains near zero ( and ). The tail can still ignore the prescribed geometry and use other hidden2 signals. Near-perfect linear-probe AUC shows that country also remains easy to read linearly.

2.5 Make the Geometry Functional and Reduce Linear Access (GFAL)

Can country remain useful while becoming harder to read as one ordinary linear direction?

I call this combined stage GFAL, for Geometry Functional and Anti-Linear. It retains the task, ClassOT, radius, and geometry-score losses; restores MixOT coverage; and adds causal geometry pressure, tail fitting, anti-linear pressure, and a squared-correlation penalty.

These additions repair different failures. Causal pressure and tail fitting make the tail respond to the three geometry channels. MixOT protects coverage, especially for the helix. Anti-linear pressure and correlation penalties weaken ordinary linear country readouts from hidden2.

2.5.1 Causal Geometry Pressure

If I replace only the geometry channels, does the model's own country logit move?

I hold complement fixed and create positive and negative edited states:

Let be the model’s tail from hidden2 to logits. The target pressure asks the tail to respond correctly: should be high while should be low. The target causal response is

I also penalize off-target spillover to avoid affecting other labels:

so that

Positive geometry should raise the country logit and negative geometry should lower it [6, 7], without moving every other output.

2.5.2 Tail Fitting

hidden2 → hidden3 → logits

Tail fitting freezes earlier layers and trains only the final tail. It teaches the tail how to decode the geometry channels into the country logit; it assumes the geometry code is already present rather than shaping it itself.

Causal geometry pressure and tail fitting are not redundant because the model can fail in two separate ways:

1. Good geometry, bad use. The three channels contain the intended structure, but the tail ignores them. Tail fitting gives the decoder a focused reason to read them.
2. Bad geometry, good reader. The tail is willing to read the channels, but the channels do not form the desired structure. Tail fitting cannot repair this; geometry losses and causal pressure do that work.

Geometry losses write the code in the right language, causal pressure checks that changing the code changes the answer, and tail fitting teaches the final reader how to read that language.

With , I use:

• Context state: . This lets the tail fit on complement values for each example, but prevents learning signal from flowing back to earlier layers to hide country information in these complement channels.
• Neutral state: , where is the mini-batch mean complement vector. If the model tail can still predict country from this state, the example-specific signal must come from the geometry channels , because the complement no longer carries example-specific clues.

The tail predicts from both states and follows the geometry score:

The model tail must predict the target label from both controlled states. As is the country logit produced by the model tail, then:

The target logit is trained to follow the intended geometry score. As should be high when is near the positive target radius and low when is near the negative target radius, tail fitting adds:

The on means this score is treated as a fixed teacher signal. The model tail must move toward the score; the geometry score itself is not adjusted to make the loss easier.

2.5.3 Anti-linear Pressure

Even after geometry contains country, the full hidden2 state can expose a simple linear country direction. I add a one-layer adversary

with loss

Gradient reversal leaves hidden2 unchanged in the forward pass but reverses its gradient in the backward pass. The adversary therefore learns to predict country, while the representation learns to make that linear prediction worse. Equivalently,

The objective is not to erase country entirely. The main task still needs country and geometry losses still require the first three channels to carry its manifold. It is to remove the arbitrary straight-line shortcut, following the gradient-reversal mechanism of Ganin et al. [8].

2.5.4 Squared Correlation Penalty

An adversary can underfit or miss one narrow channel that quietly tracks country. I therefore add a direct backup check for every hidden2 channel:

This does not prove that the complement is clean, but it closes the easy one-channel shortcut.

2.5.5 Experiment

2.5.6 Results

Table 4. Model behavior after adding causal geometry pressure, tail fitting, anti-linear pressure, and squared correlation penalty (GFAL: Geometry Functional and Anti-Linear).

GFAL makes the geometry functional. Geometry-probe AUC stays at for the sphere and for the helix, while causal delta rises from near zero to and . Geometry edits now move the model's country logit.

Linear-probe AUC falls from in the ClassOT stage to . The sphere result is cleaner; the helix retains lower coverage and higher linear access, so MixOT remains necessary.

2.6 Reduce complement shortcuts

GFAL makes geometry functional, but does not prove the remaining 61 channels are harmless.

If I remove geometry code and look only at the remaining 61 channels, can I still read country?

If yes, geometry is functional but not primary: the model has a backup route. GFAL+ adds a second one-layer adversary :

or, in probability form,

The complement-correlation penalty provides a direct backup check for any single complement channel that tracks country:

2.6.1 Why keep both anti-linear and complement adversaries?

Why do I not remove the anti-linear adversary and use only the complement adversary? Because complement cleanup is narrower: the complement adversary only looks at the  non-geometry channels.

The complement adversary asks: can country be read after removing the geometry channels?

The anti-linear adversary asks: is country still easy to read as a simple linear direction from the whole hidden2 state?

The first question is important but cannot replace the broader question from the anti-linear adversary. The whole hidden2 state contains both  and , and the tail classifier sees both. Without full-hidden2 anti-linear pressure, the model can reopen an easy linear shortcut using their mixture.

The squared correlation penalty remains for the same reason. The adversary attacks a learned linear readout, but it depends on optimization; correlation attacks the simpler failure where one hidden2 channel quietly tracks country. This stage adds complement-specific pressure without declaring the earlier shortcut solved forever. Keeping both is a guardrail: while the model cleans up , country must not become trivially exposed again in the full representation.

2.6.2 Experiment