Problem statement

The original BlueDot Technical AI Safety Puzzle 1 problem statement supplies a classifier for short English texts. It predicts eight binary features simultaneously; the features are not mutually exclusive, and the supplied model achieves more than 95% accuracy on each:

• number: contains a digit or written-out number;
• question: is phrased as a question;
• color: contains a color word;
• food: mentions food;
• sentiment: expresses positive rather than negative sentiment;
• country: contains a country name;
• person: contains a person’s name; and
• body_part: contains a body-part word.

The model is sentence-transformers/all-MiniLM-L6-v2, mean-pooled to a 384-dimensional sentence embedding, followed by four 64-dimensional ReLU hidden layers and an eight-logit multilabel output. The puzzle designates layer L as the activation after the third ReLU (hidden 2, called h2 below). Seven labels are represented linearly at L: one direction in activation space describes each label. One unknown feature F is represented differently.

The required tasks are:

1. Find F: identify which of the eight features is not represented linearly at L.
2. Explain the representation: describe the geometry used to represent F at L and show the analysis supporting that conclusion.
3. Train a stranger representation: train a model that encodes F, or another feature, in a representation more interesting than the supplied model’s representation; define and defend “more interesting.”

The supplied evidence is model.pt, 7,000 training texts, 1,500 held-out test texts, and the label ordering in feature_names.json. The requested submission is one document recording what was tried, what worked, what failed, and what structure emerged. This report is self-contained with respect to those requirements; the link above identifies the authoritative original wording and starter code, also preserved in an immutable upstream snapshot.

Summary

The unusual feature is country. At the specified hidden layer (h2), a linear probe is at chance (0.471) while a nonlinear probe reaches 0.966. The simplest high-performing geometric account is a one-dimensional interval code on the model’s food direction: the sign carries food, while the magnitude largely carries country. A raw signed projection predicts country at 0.502, whereas its absolute value and square both reach 0.946.

For Task 3, I trained a two-dimensional, unit-norm circular bottleneck on MNIST. Digit classes occupy successive angles, so even and odd digits alternate around the circle. After freezing the architecture, hyperparameters, split, and five seeds using validation data only, I evaluated once on the additional 50,000 QMNIST test examples. The mean even/odd accuracy was 0.537 for a linear probe and 0.958 for a nonlinear probe. This is a robust nonlinear advantage, though not perfect linear-probe immunity in every seed.

Task 1: Finding the unusual feature

I cached activations at the sentence embedding, each post-ReLU hidden layer, and the logits. For every feature and layer I trained:

• a logistic-regression probe; and
• a one-hidden-layer MLP probe.

Both probes were fit on the supplied training examples and scored on the supplied held-out examples. country is the only feature with a large linear/nonlinear gap at h2.

The layer trace localizes the phenomenon. Country is linearly decodable before h2, nonlinear only at h2, and linearly decodable again at h3 and the logits.

Three checks make a template artifact unlikely. Leave-one-template-out scoring gives 0.486 linear versus 0.995 nonlinear; a whitened-PCA control gives 0.503 versus 0.971; and shuffled country labels return to chance. These are probe results, not by themselves proof that the model causally uses every measured feature [1]. Here the supplied classifier’s own country output also reaches 0.964, showing that the information is behaviorally available.

Task 2: Geometry of country at h2

A magnitude code on the food direction

Let z be the standardized h2 activation and let w_food be a logistic regression direction fit to predict food on training data. I measured the scalar projection

$$p = w_{food}^{T} z.$$

The sign of p tracks food. Country changes the magnitude: country-positive examples lie nearer zero, while country-negative examples lie farther out on either side. The direction’s sign is arbitrary; the symmetric magnitude pattern is the relevant fact.

A one-dimensional decoder isolates the symmetry:

Country directions fit separately within food=0 and food=1 have cosine similarity -0.985. Thus the country direction reverses when food changes, as expected if food supplies sign and country supplies magnitude. I therefore summarize the geometry as a primarily one-dimensional, sign-symmetric interval code aligned with the food direction. The word “primarily” matters: 0.946 is strong but not exact, and this analysis does not prove that a single scalar accounts for every country-relevant computation.

This resembles the absolute-value computation studied in toy models of superposition, where opposite ReLU branches can recover a sign-symmetric quantity [2]. It does not establish why this classifier learned the code. In particular, eight binary labels in 64 dimensions do not by themselves establish capacity pressure, so I make no capacity-causation claim.

What the next layer does, and does not show

The nonlinear code does not persist to the output: country is linearly decodable at h3 (0.964) and at the logits (0.963). Two h3 neurons have large, opposite input-weight alignment with the food direction. A newly fit logistic probe on those two activations scores 0.895, but that is correlational. Using the trained model’s actual country-output weights, those two neurons alone score only 0.517. Replacing them with their training means leaves accuracy at 0.967, compared with 0.964 unablated. They carry country-correlated information, but they are neither sufficient under the trained readout nor shown to be necessary. I therefore do not call them a two-neuron decoder.

I also discard an earlier “exact rank-1 bilinear” claim. Squaring one selected projection constructs an outer-product matrix of rank one by definition. With the separately fit unstandardized direction, that decoder reaches only 0.727; it is a rank-1 quadratic approximation, not an exact account of the classifier.

Complete Tasks 1 and 2 analysis record

The headline account above was selected from the following analyses. This table includes null results and failed interpretations, not only supporting evidence. All probe fits use training data and all reported scores use the supplied 1,500 example test split unless explicitly marked as cross-validation.

Task 3: An interleaved circular representation

Construction

A small CNN maps each MNIST image to a two-dimensional bottleneck b, normalized so that $\lVert b\rVert_2=1$. For digit class $k$, a geometry loss targets angle

$$\theta_k = \frac{2\pi k}{10}$$

A linear ten-class head predicts the digit [4]. The derived binary feature is parity: even digits are classes 0, 2, 4, 6, 8, and odd digits are 1, 3, 5, 7, 9. They alternate around the unit circle, so parity occupies five disjoint angular regions rather than one half-space. Unit normalization is essential because it prevents radius from becoming a linear shortcut.

Frozen evaluation protocol

The original exploratory run repeatedly inspected a random split of all 70,000 MNIST examples. Because that split included canonical MNIST test images, I treat all results from it as contaminated and do not report them as held-out evidence.

The corrected protocol was fixed before the confirmatory audit:

• canonical MNIST training partition only;
• stratified 50,000/10,000 train/validation split, split seed 17291;
• training seeds 11, 29, 47, 71, 101;
• 30 epochs, Adam learning rate 0.001, batch size 256;
• geometry-loss weight 1.0;
• fixed epoch count, with no checkpoint or seed selection on audit outcomes; and
• a single confirmatory evaluation on QMNIST test50k, the 50,000 additional reconstructed test digits that do not duplicate the standard MNIST test set [3].

Training code never loads QMNIST. The audit script refuses to run if its result CSV already exists. Its first invocation stopped before predictions or metrics because raw QMNIST targets expose eight metadata columns; class-column extraction was unit-tested, then the unchanged frozen checkpoints were evaluated once. Checkpoint SHA-256 hashes and the event are recorded in the audit artifacts.

Confirmatory results

The intervals summarize variation across five training seeds as $\bar{x} \pm t_{0.975,4}s/\sqrt{5}$ with $t_{0.975,4}=2.776$; they are not per-example binomial intervals.

The mean nonlinear-minus-linear parity gap is 0.421; every seed has a gap of at least 0.374. The nonlinear result is consistently high. The linear result is near the majority baseline on average, but seed 71 reaches 0.601, so the defensible claim is strong and repeatable nonlinear advantage from circular interleaving, not universal chance-level linear probing. Direct-head digit accuracy is also seed-sensitive (0.795 to 0.943), while a refit linear digit probe averages 0.949; this is a limitation of optimization in the frozen run, not something removed after seeing audit data.

Complete Task 3 experiment record

The circular MNIST model above was the final confirmatory experiment. For completeness, the table below records every preceding and subsequent Task 3 family, including failures. Except for Idea 6’s frozen MNIST/QMNIST protocol, Ideas 1-5 and 7-12 belong to architecture search: they repeatedly used the supplied test activations, usually with one seed. Their numbers are useful for explaining what was tried and what geometry appeared, but they are test-contaminated exploratory evidence rather than independent held-out results. No run was promoted or omitted because of its test score.

The sequence exposes a consistent failure mode. A binary label can satisfy a nominally circular, helical, harmonic, or ring-shaped objective by occupying only one or two easy regions. Those regions are usually linearly separable, or the feature disappears entirely. The confirmatory MNIST construction avoids this by using ten supervised classes to populate the full unit circle before deriving a binary parity label. The d=4 bottleneck is the most interesting emergent text result, but it hides person, not country, and trades away performance on some other labels.

Diagnostics from all exploratory geometries

The plots below are the checked-in diagnostics for every Task 3 family not already shown in Figure 3. They are included even when the intended structure did not emerge; captions report the failure rather than the training target.

Conclusion

Tasks 1 and 2 identify a localized nonlinear representation: country is hidden from linear readout at h2 by a sign-symmetric magnitude code tied to the food direction, then relinearized at h3. The causal evidence does not support a specific two-neuron decoder, an exact rank-1 mechanism, or a capacity-based origin story.

Task 3 shows a more elaborate binary representation by embedding a ten-class variable on a unit circle and reading parity from alternating arcs. On an additional QMNIST audit set and across five frozen seeds, nonlinear parity probes consistently outperform linear probes by a large margin. The seed variation is part of the result: the construction is strongly probe-resistant on average, but not perfectly immune to linear leakage.

Reproducibility

The hosted report, complete code, and checked-in artifacts are in the public submission repository, a fork of the original puzzle repository. Clone the submission branch, then explicitly fetch and fast-forward it before running any command below:

git clone --branch release/puzzle-1-submission --single-branch \
   https://github.com/janmenjayap/bluedot-tais-puzzle.git
cd bluedot-tais-puzzle
git fetch origin release/puzzle-1-submission
git merge --ff-only FETCH_HEAD
./setup.sh
conda activate bluedot-impact-puzzle-1-py311
pytest -q

setup.sh creates a Python 3.11 Conda environment, installs the version-bounded dependencies in requirements.txt (including Torch 2.2 and torchvision 0.17), and runs pip check. It does not alter or checksum the supplied model and text data. pytest -q verifies activation taps, probe utilities, Task 3 models, and the frozen MNIST/QMNIST protocol.

The branch contains the latest documentation. The submitted analysis, scripts, and artifacts are also archived in the immutable v1.0.0 release at commit 341d3d5. To reproduce that exact snapshot after cloning, run:

git fetch origin tag v1.0.0
git switch --detach v1.0.0

Verify the published evidence

The report can be checked without retraining. The first command verifies the code; the next two display the frozen five-seed result and the hashes recorded for the exact checkpoints used in the one-time audit:

pytest -q
cat artifacts/results/19_mnist_circular_audit_summary.csv
cut -d, -f1-2 artifacts/results/19_mnist_circular_audit_seeds.csv
shasum -a 256 artifacts/results/18_mnist_circular_seed_*.pt

The hashes printed by shasum should match the checkpoint_sha256 column after accounting for seed order. Protocol metadata and the pre-outcome loader event are in 18_mnist_circular_protocol.json and 19_mnist_circular_audit_record.json.

The supplied puzzle inputs used here have these SHA-256 hashes:

The checked-in HTML is self-contained. With Pandoc 3.9, regenerate it from the auditable Markdown source and stylesheet as follows:

pandoc docs/writeup.md --standalone --toc --embed-resources --mathml \
   --syntax-highlighting=none --css docs/submission.css --output docs/report.html

Reproduce Tasks 1 and 2

Run the analyses in dependency order from the repository root. These commands rebuild the activation cache and the complete Tasks 1-2 evidence table:

for script in \
   00_cache_activations \
   01_sanity_check \
   02_linear_probes \
   03_nonlinear_probes \
   04a_significance \
   04b_template \
   04c_controls \
   04d_layer_trace \
   05_geometry \
   06_country_food \
   07_mechanism \
   07b_h3_circuit
do
   python "scripts/${script}.py"
done

The generated CSVs and figures are artifacts/results/01_* through artifacts/results/07b_*; artifacts/activations/acts.npz is the shared cache.

Independently rerun the confirmatory Task 3 audit

Task 3’s frozen training entry point is scripts/18_train_mnist_circular.py; it loads only the canonical MNIST training partition. The audit entry point is scripts/19_analyze_mnist_circular.py; it loads QMNIST test50k and deliberately refuses to run while a completed audit CSV exists. In a disposable reproduction clone, preserve the published artifacts before invoking the guarded path:

mkdir -p artifacts/published-mnist-audit
mv artifacts/results/18_mnist_circular_* \
    artifacts/results/19_mnist_circular_* \
    artifacts/published-mnist-audit/
python scripts/18_train_mnist_circular.py
python scripts/19_analyze_mnist_circular.py
diff -u artifacts/published-mnist-audit/19_mnist_circular_audit_summary.csv \
            artifacts/results/19_mnist_circular_audit_summary.csv

The frozen split, seeds, epoch count, learning rate, batch size, and geometry weight are constants in src/puzzle/mnist_protocol.py. Python, Torch, torchvision, and device metadata are written into the protocol JSON. Random seeds are fixed, but exact checkpoint bytes can remain backend-dependent; the published SHA-256 values identify the checkpoints supporting this report.

Rerun the historical Task 3 archive

For completeness, the remaining Task 3 scripts run in the order below. They train multiple models and can take substantially longer than the Tasks 1-2 analysis. Run them only in a disposable clone because they overwrite historical artifacts. Their use of repeatedly inspected test data means rerunning them does not convert their results into confirmatory evidence.

for script in \
   08_train_xor 09_analyze_xor \
   10_train_rotation 11_analyze_rotation \
   12_train_jepa 13_analyze_jepa \
   14_train_helix 15_analyze_helix \
   16_train_superposition 17_analyze_superposition \
   20_train_sae_h2 21_analyze_sae_h2 \
   22_cpd_puzzle_analytic \
   23_train_bilinear_xor 24_analyze_bilinear_xor \
   25_train_bilinear_text 26_analyze_bilinear_cpd \
   27_train_bottleneck 28_analyze_bottleneck \
   29_train_squarewave 30_analyze_squarewave \
   31_train_fourier_comb 32_analyze_fourier_comb \
   33_train_linked_rings 34_analyze_linked_rings
do
   python "scripts/${script}.py"
done

The script-to-experiment mapping is the numbered Task 3 ledger above. The authoritative evidence classification is also recorded in scripts/README.md: 18/19 are confirmatory, 00-07b analyze the supplied fixed model under its intended train/test split, and the other Task 3 runs are historical exploration.

References

1. Alain, G., and Bengio, Y. (2016, revised 2018). “Understanding intermediate layers using linear classifier probes.” arXiv:1610.01644. https://arxiv.org/abs/1610.01644
2. Elhage, N., et al. (2022). “Toy Models of Superposition.” Transformer Circuits Thread. https://transformer-circuits.pub/2022/toy_model/index.html
3. Yadav, C., and Bottou, L. (2019). “Cold Case: The Lost MNIST Digits.” Advances in Neural Information Processing Systems 32. https://proceedings.neurips.cc/paper/2019/hash/51c68dc084cb0b8467eafad1330bce66-Abstract.html
4. LeCun, Y., Bottou, L., Bengio, Y., and Haffner, P. (1998). “Gradient-Based Learning Applied to Document Recognition.” Proceedings of the IEEE, 86(11), 2278-2324. https://doi.org/10.1109/5.726791