{"slug": "can-we-teach-a-model-to-encode-a-semantic-feature-on-a-chosen-manifold-in-just", "title": "Can we teach a model to encode a semantic feature on a chosen manifold in just three channels?", "summary": "A researcher received an Honorable Mention in BlueDot's Technical AI Safety Puzzle #1 for training a small MLP to encode a country feature on a chosen nonlinear manifold in three reserved channels, demonstrating that a model can learn to use a predefined geometric representation while solving the original classification task.", "body_md": "This is my submission to [BlueDot's Technical AI Safety Puzzle #1](https://bluedot.org/puzzles/technical-ai-safety), for which I received an **Honorable Mention**.\n\nCongratulations to Gustavo Korzune Gurgel, Patryk Perduta ([his amazing write-up](https://www.lesswrong.com/posts/TRno3dELwmMHYYH5v/bluedot-technical-ai-safety-puzzle-submission-that-got-me-a)), Sam Spilllard, Karine Levonyan, and Michael Zlatin for their recognition in the puzzle.\n\nMy 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.*\n\nMy Task 1 and 2 write-up is available on [my homepage](https://phusroyal.github.io/writing/bluedot-1st-puzzle-q12/), and the [interactive/more intuitive version of this article](https://phusroyal.github.io/writing/bluedot-1st-puzzle-q3/).\n\nI welcome discussion, feedback, and collaborations that could extend this idea.\n\nYou can check out the code for this article in my [GitHub repository](https://github.com/phusroyal/Phu-BlueDot_1st_puzzle).\n\n[BlueDot's Technical AI Safety Puzzle #1](https://bluedot.org/puzzles/technical-ai-safety) 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:\n\nThis post addresses the third task. I train a new five-layer MLP and constrain `country`\n\nto use a chosen three-dimensional manifold while testing whether the classifier relies on that code.\n\nThe guiding question is:\n\nCan we choose a representation geometry first, then train the model so that the feature follows that geometry while still solving the original task?\n\nThis 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.\n\nThe manifold must not be decorative. I therefore ask both:\n\nDid the hidden activations form the desired shape?\n\nBut also:\n\nDid the model actually use that shape to solve the task?\n\nMy design has five requirements:\n\nThe first stage has no model. I only define what the model will later be asked to learn.\n\nFor each example, the dataset gives the binary label . It says whether `country`\n\nis present, but does not say where the example belongs on a desired manifold.\n\nFor a helix, supplies no position , radius , or tube angle . If those coordinates were observed, I could use point-wise supervision to train my model.\n\nInstead, I use a distributional target:\n\nHere, 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].\n\nI 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.\n\nLet be a direction sampled uniformly from the unit sphere , the set of points in 3-dimensional space with Euclidean norm 1.\n\nLet be a radius. The raw sphere point is\n\nThe positive class uses a smaller radius than the negative class:\n\nwith 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`\n\nis encoded by an inner versus outer shell, not by direction.\n\nLet be a position along the helix and its vertical pitch (to control how fast the helix rises). The helix core curve is\n\nTwo local cross-section directions around that center line are\n\nWith cross-section angle and tube radius , a point in the tube is\n\nThe class-conditional distributions are\n\nI use , , and . Both classes trace the same curved scaffold.\n\nI tried three turns, but the tail layer did not reliably read them under the remaining constraints; one turn was learnable, usable, and testable.\n\nThe 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,\n\nFor the sphere shell,\n\nand for the helix tube,\n\nHere 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.\n\nBefore training, I draw normalized target anchors and .\n\nEach text input is encoded by `sentence-transformers/all-MiniLM-L6-v2`\n\n, then mean-pooled into . A five-layer ReLU MLP maps it to eight logits and is trained with binary cross-entropy:\n\nThe 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.\n\nI 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.\n\nBecause hidden2 is post-ReLU, I add positive offset so a signed manifold can live in activation space:\n\nLet 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:\n\nFor the sphere bottleneck,\n\nwhere is the azimuth angle , is the vertical coordinate, and has maximum allowed radius .\n\nFor the helix bottleneck,\n\nwhere , , and .\n\nAt this stage, I optimize only . The purpose is not yet to force good class geometry. I ask only:\n\nCan the classifier survive this architectural constraint?\n\nAlongside mean and country AUC, I use four analyses:\n\nIts ROC AUC shows whether the reserved code contains the label in the intended geometric form; it does not prove the tail uses that code.\n\n*Table 1. Model behavior after adding the geometry bottleneck. Higher is better unless marked ↓.*\n\nGeometry | Mean task AUC ↑ | Country AUC ↑ | Geometry probe AUC ↑ | Linear probe AUC ↓ | Coverage entropy ↑ | Causal delta ↑ |\n|---|---|---|---|---|---|---|\nSphere shell | 0.9975 | 0.9995 | 0.2021 | 0.9996 | 0.7971 | 0.0187 |\nHelix tube | 0.9963 | 0.9997 | 0.1182 | 0.9997 | 0.6561 | 0.0357 |\n\nThe 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.\n\nCan the learned geometry codes cover the target geometry instead of collapsing to a small region?\n\nThe 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.\n\nTo prevent this collapse, I add **unconditional mixture optimal transport (MixOT)**, which spreads the unlabeled mini-batch across the overall target geometry:\n\nwhere is the country-positive rate.\n\nFor mini-batch , let be its n learned 3D geometry codes and sample equally many anchors from .\n\nThink 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.\n\nTherefore, for a learned code and a target anchor . The pairwise transport cost is\n\nEntropic OT solves\n\nsubject to\n\nHere 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]\n\nTo test whether class-conditional geometry works on its own, I replace the previous MixOT term with:\n\nI add Radius MAE, the mean absolute error between learned and desired class radius (lower is better).\n\n*Table 2. Model behavior after adding unconditional mixture optimal transport.*\n\nGeometry | Stage | Mean task AUC ↑ | Country AUC ↑ | Geometry probe AUC ↑ | Linear probe AUC ↓ | Coverage entropy ↑ | Causal delta ↑ | Radius MAE ↓ |\n|---|---|---|---|---|---|---|---|---|\n| Bottleneck |\n| | 0.2021 | 0.9996 | 0.7971 | 0.0187 | 0.7243 |\nMixOT | 0.9974 | 0.9990 |\n|\n|\n|\n|\n| |\n| Bottleneck | 0.9963 | | 0.1182 | 0.9997 | 0.6561 | 0.0357 | 0.6653 |\nMixOT |\n| 0.9993 |\n|\n|\n|\n|\n|\n\nMixOT makes coverage almost perfect, but geometry-probe AUC remains weak, radius MAE remains large, and causal delta remains nearly zero.\n\nThis is expected: it spreads the unlabeled batch but does not assign positives to and negatives to .\n\nPut positive examples in the positive part of the shape, and negative examples in the negative part.\n\nMixOT 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.\n\nFor , let and sample anchors from . Then\n\nwhere is the number of included label groups; I skip a class with fewer than two examples in the batch.\n\nClassOT 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.\n\nLet be distance from the origin for the sphere or distance to the nearest helix core point for the helix:\n\nGeometry-score loss asks whether the geometry itself would classify a point correctly. Define\n\nand the logit-like score\n\nIt is high near the positive radius and low near the negative radius, so I optimize\n\nClassOT 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.\n\nFor 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 .\n\n*I remove the previous stage's mixture-OT term to test whether class-conditional geometry works on its own.*\n\n*Table 3. Model behavior after replacing unconditional mixture OT with class-conditional OT.*\n\nGeometry | Stage | Mean task AUC ↑ | Country AUC ↑ | Geometry probe AUC ↑ | Linear probe AUC ↓ | Coverage entropy ↑ | Causal delta ↑ | Radius MAE ↓ |\n|---|---|---|---|---|---|---|---|---|\nSphere shell | Bottleneck |\n| 0.9995 | 0.2021 | 0.9996 | 0.7971 | 0.0187 | 0.7243 |\nMixOT | 0.9974 | 0.9990 | 0.5585 |\n|\n|\n| 0.3012 | |\nClassOT | 0.9962 |\n|\n| 0.9998 | 0.9970 | 0.0188 |\n| |\nHelix tube | Bottleneck | 0.9963 |\n| 0.1182 | 0.9997 | 0.6561 | 0.0357 | 0.6653 |\nMixOT |\n| 0.9993 | 0.3095 |\n| 0.9942 | 0.0368 | 0.5582 | |\nClassOT | 0.9949 | 0.9996 |\n| 0.9997 |\n|\n|\n|\n\nClassOT 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.\n\nCan country remain useful while becoming harder to read as one ordinary linear direction?\n\nI 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.\n\nThese 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.\n\nIf I replace only the geometry channels, does the model's own country logit move?\n\nI hold complement fixed and create positive and negative edited states:\n\nLet 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\n\nI also penalize off-target spillover to avoid affecting other labels:\n\nso that\n\nPositive geometry should raise the country logit and negative geometry should lower it [6, 7], without moving every other output.\n\nhidden2 → hidden3 → logits\n\nTail 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.\n\nCausal geometry pressure and tail fitting are not redundant because the model can fail in two separate ways:\n\nGeometry 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.\n\nWith , I use:\n\nThe tail predicts from both states and follows the geometry score:\n\nThe model tail must predict the target label from both controlled states. As is the country logit produced by the model tail, then:\n\nThe 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:\n\nThe 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.\n\nEven after geometry contains country, the full hidden2 state can expose a simple linear country direction. I add a one-layer adversary\n\nwith loss\n\nGradient 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,\n\nThe 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].\n\nAn adversary can underfit or miss one narrow channel that quietly tracks country. I therefore add a direct backup check for every hidden2 channel:\n\nThis does not prove that the complement is clean, but it closes the easy one-channel shortcut.\n\n*Table 4. Model behavior after adding causal geometry pressure, tail fitting, anti-linear pressure, and squared correlation penalty (GFAL: Geometry Functional and Anti-Linear).*\n\nGeometry | Stage | Mean task AUC ↑ | Country AUC ↑ | Geometry probe AUC ↑ | Linear probe AUC ↓ | Coverage entropy ↑ | Causal delta ↑ | Radius MAE ↓ |\n|---|---|---|---|---|---|---|---|---|\nSphere shell | Bottleneck |\n| 0.9995 | 0.2021 | 0.9996 | 0.7971 | 0.0187 | 0.7243 |\nMixOT | 0.9974 | 0.9990 | 0.5585 | 0.9993 | 0.9991 | 0.0206 | 0.3012 | |\nClassOT | 0.9962 |\n|\n| 0.9998 | 0.9970 | 0.0188 |\n| |\nGFAL | 0.9799 | 0.9995 | 0.9997 |\n|\n|\n| 0.1444 | |\nHelix tube | Bottleneck | 0.9963 |\n| 0.1182 | 0.9997 | 0.6561 | 0.0357 | 0.6653 |\nMixOT |\n| 0.9993 | 0.3095 | 0.9994 | 0.9942 | 0.0368 | 0.5582 | |\nClassOT | 0.9949 | 0.9996 |\n| 0.9997 |\n| 0.0389 |\n| |\nGFAL | 0.9739 | 0.9957 | 0.9940 |\n| 0.7510 |\n| 0.1887 |\n\nGFAL 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.\n\nLinear-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.\n\nGFAL makes geometry functional, but does not prove the remaining 61 channels are harmless.\n\nIf I remove geometry code and look only at the remaining 61 channels, can I still read country?\n\nIf yes, geometry is functional but not primary: the model has a backup route. GFAL+ adds a second one-layer adversary :\n\nor, in probability form,\n\nThe complement-correlation penalty provides a direct backup check for any single complement channel that tracks country:\n\nWhy 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.\n\nThe complement adversary asks:\n\ncan country be read after removing the geometry channels?\n\nThe anti-linear adversary asks: *is country still easy to read as a simple linear direction from the whole hidden2 state?*\n\nThe 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.\n\nThe 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.\n\n*Table 5. Model behavior after adding complement-adversary pressure (GFAL+).*\n\nGeometry | Stage | Mean task AUC ↑ | Country AUC ↑ | Geometry probe AUC ↑ | Linear probe AUC ↓ | Coverage entropy↑ | Causal delta ↑ | Radius MAE↓ | Complement AUC ↓ |\n|---|---|---|---|---|---|---|---|---|---|\nSphere shell | Bottleneck |\n| 0.9995 | 0.2021 | 0.9996 | 0.7971 | 0.0187 | 0.7243 | 0.9996 |\nMixOT | 0.9974 | 0.9990 | 0.5585 | 0.9993 | 0.9991 | 0.0206 | 0.3012 | 0.9981 | |\nClassOT | 0.9962 | 0.9996 | | 0.9998 | 0.9970 | 0.0188 | 0.1286 | 0.9998 | |\nGFAL | 0.9799 | 0.9995 | 0.9997 |\n| 0.9993 | 3.0845 | 0.1444 | 0.5758 | |\nGFAL+ | 0.9674 |\n| 0.9997 | 0.6204 |\n|\n|\n|\n| |\nHelix tube | Bottleneck | 0.9963 |\n| 0.1182 | 0.9997 | 0.6561 | 0.0357 | 0.6653 | 0.9997 |\nMixOT |\n| 0.9993 | 0.3095 | 0.9994 | 0.9942 | 0.0368 | 0.5582 | 0.9994 | |\nClassOT | 0.9949 | 0.9996 |\n| 0.9997 |\n| 0.0389 |\n| 0.9996 | |\nGFAL | 0.9739 | 0.9957 | 0.9940 | 0.6710 | 0.7510 | 1.8980 | 0.1887 | 0.6071 | |\nGFAL+ | 0.9576 | 0.9949 | 0.9927 |\n| 0.7787 |\n| 0.2111 |\n|\n\nGFAL+ preserves task performance, strong geometry probes, and causal geometry for both manifolds. Its complement result is asymmetric: complement AUC falls from to for the helix but rises from to for the sphere. Thus it reduces a helix backup route but does not establish perfect information isolation for either geometry.\n\n*If I edit only the geometry channels after training, does the trained classifier actually follow that edit?*\n\nThe earlier probes show that geometry channels contain country information; they do not show that the classifier uses them. I freeze the model, select a balanced held-out subset, and edit only hidden2 geometry channels. I use four interventions:\n\nFor replacement and swap, I also measure **specificity**: whether country moves much more than the other seven logits.\n\nI replace with zero, then compute target AUC. The ablation drop is\n\nIf I remove the geometry signal, does country prediction get worse?\n\nI replace learned geometry with positive or negative anchors while preserving the original complement:\n\nHere is the intervened hidden2 state. Because is copied from the original example, the replacement geometry is the only intended cause of output movement. The causal target delta is\n\nFor positive-negative pairs , I exchange only their geometry channels. Let and be the selected positive and negative indices. The original hidden states are:\n\nNow swap only the geometry channels:\n\nThe swap target shift is\n\nIt is large when negative geometry lowers positive examples and positive geometry raises negative examples.\n\nI choose a deterministic path through each target geometry from the positive to the negative radius:\n\nRather than jumping directly between the endpoints, I traverse intermediate geometry points and check whether the country logit changes smoothly. For the sphere, I fix one direction and increase radius from to . For the helix, I fix one core position and tube angle, then increase only the tube radius. I hold the complement at mean :\n\nLet is the country logit at step . Since the path moves from positive geometry to negative geometry, the expected behavior is that country logit decreases. The path target delta is:\n\nThe path monotonic fraction is:\n\nmeans every adjacent step moves in the expected direction.\n\nSpecificity compares country movement with average off-target movement:\n\nA high ratio means the geometry edit mainly affects country rather than every feature.\n\n*Table 6. Causal-use validation on the trained model from the previous stage.*\n\nGeometry | Ablation AUC drop | Ablation target AUC | Causal delta | Swap target shift | Swap specificity | Path target delta | Path monotonic fraction | Specificity ratio |\n|---|---|---|---|---|---|---|---|---|\nSphere shell | 0.4675 | 0.5317 | 3.5004 | 4.4433 | 92.0223 | 3.4963 | 1.0000 | 73.4252 |\nHelix tube | 0.4279 | 0.5654 | 2.6884 | 7.1863 | 128.0430 | 5.1727 | 1.0000 | 47.3151 |\n\nBoth geometries pass every causal check. Ablation nearly removes country prediction, leaving both near chance. Replacement moves country logits in the expected direction with specificity ratios and . Each path has monotonic fraction .\n\nThe swap test is the strongest learned-code check because it exchanges the model's own geometry codes rather than synthetic anchors. It yields target shifts for the sphere and for the helix, with specificities and . The learned geometry is therefore not merely visually aligned with country: moving it between examples changes the model's prediction path as expected.\n\nThe sphere is geometrically cleaner, with higher coverage entropy and lower radius MAE, but both geometries show strong causal use. These tests establish a causally active pathway, not perfect information isolation: the complement can remain an alternative country-information source.\n\nThis experiment shows that a small MLP can encode a semantic feature through a pre-chosen nonlinear manifold in only three channels while preserving task performance. A sphere shell and helix tube both support a code that is geometrically well formed and causally used by the classifier.\n\nThe result is not that `country`\n\nhas been isolated perfectly. GFAL+ leaves evidence of complement leakage, and the helix is a harder geometry to maintain than the sphere. The useful conclusion is narrower: geometry can be specified before training, made readable from a compact code, and tested causally rather than accepted because it looks interpretable.", "url": "https://wpnews.pro/news/can-we-teach-a-model-to-encode-a-semantic-feature-on-a-chosen-manifold-in-just", "canonical_source": "https://www.lesswrong.com/posts/ZwEer94AefjdW4933/can-we-teach-a-model-to-encode-a-semantic-feature-on-a", "published_at": "2026-07-27 11:44:06+00:00", "updated_at": "2026-07-27 12:01:27.690694+00:00", "lang": "en", "topics": ["ai-safety", "machine-learning", "neural-networks", "ai-research"], "entities": ["BlueDot", "Gustavo Korzune Gurgel", "Patryk Perduta", "Sam Spilllard", "Karine Levonyan", "Michael Zlatin"], "alternates": {"html": "https://wpnews.pro/news/can-we-teach-a-model-to-encode-a-semantic-feature-on-a-chosen-manifold-in-just", "markdown": "https://wpnews.pro/news/can-we-teach-a-model-to-encode-a-semantic-feature-on-a-chosen-manifold-in-just.md", "text": "https://wpnews.pro/news/can-we-teach-a-model-to-encode-a-semantic-feature-on-a-chosen-manifold-in-just.txt", "jsonld": "https://wpnews.pro/news/can-we-teach-a-model-to-encode-a-semantic-feature-on-a-chosen-manifold-in-just.jsonld"}}