Below is a summary of my solution to BlueDot’s TAIS Puzzle #1, you can check my blog for the full write-up. The puzzle gave you a five-layer MLP on top of frozen MiniLM embeddings, predicting eight binary features. Seven features are linearly recoverable at layer h2, but one isn’t. The task was to find it, describe the geometry, and then engineer something weirder. The emphasis of this post is on two parts we found most interesting: establishing causal proof via Distributed Alignment Search (DAS), and a blueprint for constructing hidden, slice-proof representation geometries, what we called a velocity ring.

Key Findings

• Found the hidden feature country, which collapses to random chance at layer h2because it is hidden as a pure XOR with food.
• Applied causal verification via DAS to isolate a 1-D non-linear subspace at h2. Activation patching this single coordinate flips downstream model outputs, which proves the geometry is actively causal.
• Constructed a novel, continuous, zero-mean ring geometry hidden entirely inside the inter-layer velocity update, achieving absolute linear blindness and defeating single-feature slicing. (Git repo for the weirder rep)

Tasks 1 & 2. Which is the nonlinear feature and Why?

Standard linear probes see nothing at layer h2 because positive country instances sit directly "inside" negative instances along a single axis. For each layer × feature pair, we fit a linear logistic-regression probe on training activations and score it on test activations. A drop in linear probe accuracy signals a non-linear representation.

We see that the Country loses accuracy to 0.43 at h2 (blue block) and then recovers at h3. Every other feature stays above 0.97 at h2. This suggests a nonlinear representation for the Country feature. A non-linear MLP probe at h2 recovers the Country to 0.96. So the information is present at h2, but it is not linear.

Below on the histograms and PCA plots we see again that no linear probe reads Country at h2 because the positive and negative distributions share the same mean. The positive examples sit sandwiched inside the negatives along a single axis:

Figure 2. Histogram of each feature by their binary label.

Figure 3. PCA distribution of each feature activations at layer-2.

Finally, as an additional test we try to see if we can restore linearity by projection-rotation manipulation.  We see a classical nested layout where country=1 is wrapped inside country=0 at level h2, while linear at layers h1 and h3:

Figure. Best linear-probe score for country at h1, h2, h3: separated at h1 (0.99), collapsed at h2 (0.43), recovered at h3 (0.96). The puzzle folds the country signal into an XOR at h2 and unfolds it one layer later. This characterizes the model's representation dynamics but not how the weights were trained to produce them.

Now what we do is, first, we project the h2 into the within-food country axis, and then multiply by the food sign (the XOR-undo). That makes the country linearly separable again, increasing the probe accuracy from 0,43 all the way to 0.93 (!).

Figure Left: at h2, country=1 (red) is sandwiched inside country=0 (gray) — no line separates them. Middle: marginalising over food, the two classes overlap. Right: the degree-2 fold (projection × food-sign) makes country linearly separable, recovered by a plane at 0.93. The minimal nonlinearity is a single product with food, which is exactly the XOR.

Task 3 for Weirder Representation

We have tried training a more weird representation for the feature Sentiment. We introduce a dynamic encoding where information is stored purely in the transition between layers, rather than the state at any single layer. This representation is invisible to standard linear interpretability tools at every step of the computation, yet yields >90% accuracy in the final output.

Specifically, we encode a feature with a different, continuous geometry: a zero-mean ring read by its radius. We place it in the velocity, the inter-layer update v=h3-h2. To keep the rest of the network identical to the original, we keep the full puzzle head (l1...l5) at full depth for the other seven features and host the ring in a reserved 2-D sub-space of the real update that the readout for the others never touches. In that sub-space the feature is linearly invisible and no single-feature slice unfolds it.

The ideas is to write the update between hidden layer 2 and 3 as a velocity v = h3 − h2. We force Sentiment to be carried only by the geometry of this update in two reserved dimensions: positive examples are pushed onto a circle of radius 1, negatives sit at the origin, and the positives' angle is made independent of every label so they cover the whole circle.

The result is a ring whose mean is zero: so no linear probe and no PCA direction separates the classes in that sub-space. Yet its radius cleanly decodes sentiment. Because separation needs the radius (a non-linear scalar), conditioning on any other single feature cannot unfold it.

Let's look at the architecture change. The head is the original puzzle MLP kept at full depth: a frozen all-MiniLM-L6-v2 encoder (mean-pooled 384-dim x) feeds l1...l5 exactly as in the puzzle, with post-ReLU states h2 and h3. We make three specific changes related to how the Sentiment feature is read and how its channel is reserved.

Two properties follow by construction. The rank makes the angles uniform on the circle, so positives spread evenly, and the ring's mean is almost 0. And a classifier can separate two clouds only along a direction where their means differ, so a zero-mean ring has no separating hyperplane (linear probes and PCA see nothing).

Figure 3.4 The reserved dims of the genuine update v=h3-h2: positive examples (red) lie on a zero-mean circle (radius ≈ 0.87), negatives (gray) at the origin (radius ≈ 0.02) (left). The radius separates the classes cleanly (middle) while a linear probe on the reserved dims is at chance (right). The full 64-d update, by contrast, still carries sentiment linearly.

Figure 3.5 Linear-probe accuracy by feature and channel (h0, h1, h2, h3, linear @ reserved dims of the update, non-linear @ reserved dims, model output). Every feature is linearly readable through h0->h3; the reserved dims of the update are linearly blind for all features, and only sentiment is recovered there non-linearly.

We tried the same trick in a shared state, but it failed.

Is the velocity needed, or could we put the same ring directly in a layer state? We tried a zero-mean ring with a radius readout, now placed in two dimensions of the shared 64-dim state h2, with strong mean-matching to drive sentiment's linear shadow out of the rest of the state. The ring forms and its radius decodes sentiment (0.98), but the full shared state stays linearly readable (0.93): the linear shadow does not leave.

This is the same wall as reconstructing the puzzle (Section 4): a feature that a state must share with the others is redundantly encoded across its dimensions, so mean-matching leaves a (standardised-probe-)detectable residual, and pushing harder with an adversary destroys the whole state.

The reserved-sub-space velocity ring escapes this because the two ring dims are reserved. Specifically,  l5 is masked to zero on them, so no other feature's readout shares those coordinates. The ring therefore never has to coexist with another feature's linear encoding in the same dimensions, and a probe restricted to the reserved sub-space sits at chance (0.32).

**See Also:** If you are exploring alternative approaches to Task 3, definitely check out Girish Gupta's elegant solution. While our Velocity Ring focuses on isolating features within the inter-layer update trajectory, Girish cleverly smuggles an entire 3D globe into the model's shared hidden layer by mapping country activations directly to their real-world latitude and longitude coordinates. It is an incredibly creative look at what can be hidden inside network weights.