Can we teach a model to encode a semantic feature on a chosen manifold in just three channels?

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 Model architecture provided with the puzzle. The investigated representation is the output of the third ReLU. 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: Ca…