ThetaScan v0.1: open scan-parallel nonlinear memory with preliminary 17M LM results

I ran two small diagnostics against the public implementation. These were not 17M training replications and did not use the unpublished research checkpoint.

1. Threshold sanity check

With matched weights:

• a zero threshold reproduced ordinary ReLU² exactly in the test;
• all threshold entries received nonzero gradients;
• setting a positive threshold materially changed feature activity and hidden outputs.

For example, in one controlled sample:

The mean absolute hidden-feature change was about 0.17.

This does not show that the learned benchmark thresholds reached 0.25, or that thresholding caused the BPB difference. It only shows that this is a real trainable mechanism axis rather than inert bookkeeping.

2. Public fixed-map structure

I also inspected the key-derived sign maps at the public benchmark dimensions. Across the 32 per-head W1/W2 maps examined:

• every map had maximum possible rank;
• the median condition number was about 5.56;
• the maximum was about 5.94;
• no exact duplicate or sign-reversed rows were found;
• row-coherence tails were moderate rather than pathological.

So I did not find an obvious rank-collapse, duplicate-row, or extreme-conditioning problem in the tested maps.

3. Expansion was not equivalent to duplicating slots

At random initialization, I compared:

• the base 192-feature representation;
• a 384-slot control made by simply duplicating those features;
• the public fixed-sign expansion to 384 effective features;
• an independent 384-row reference.

The fixed expansion therefore produced additional nonlinear feature geometry; it was not merely counting the same 192 features twice. This is a positive structural sanity check for the expansion mechanism.

It is still not evidence that expansion caused the reported trained-model improvement.

4. Thresholding changed geometry, but geometry proxies were not sufficient

Applying artificial thresholds to the expanded random-initialized features increased sparsity and effective rank:

However, stronger thresholding did not monotonically improve address identification under load.

In a random-initialized address-overlap diagnostic with moderately correlated keys, unthresholded expansion often improved target margins relative to dense features, while stronger thresholds sometimes reduced top-1 recovery at higher loads.

This suggests that the following should not be used as standalone capacity proxies:

• higher effective rank;
• lower average or tail cosine;
• greater sparsity;
• smaller denominator mass.

A trained-model diagnostic should pair them with task-facing quantities such as:

• target-versus-strongest-distractor margin;
• target top-1 rate;
• target share of the normalized denominator;
• effective support size;
• feature-mass concentration;
• behavior under rewrite and hard-negative collisions.

Again, the thresholds above were manually chosen and the test used random projections. It does not establish whether learned thresholding helps or hurts the trained language model.

5. Portable backend parity

On CPU float32, I compared the public naive, quad, and cumsum paths over dense, normalized, and expanded-thresholded configurations, including non-power-of-two lengths.

Across the tested cases:

• maximum forward absolute difference: about 1.43e-6;
• maximum input-gradient difference: about 6.52e-9;
• maximum parameter-gradient difference: about 1.19e-7;
• minimum gradient cosine: above 0.9999998;
• targeted tests: 28 passed, 1 skipped.

This is useful positive evidence for the current portable implementation. It weakens the hypothesis that the published result is explained by a large disagreement among those reference paths.

It does not test BF16, CUDA, FLA, a future fused kernel, very long sequences, or wall-clock scaling.

I think it would help future readers if the result table eventually distinguished two kinds of evidence.

Best-preset comparison

This asks:

Under the declared parameter allocation and each arm’s selected training policy, which complete configuration performed best?

The current 7,500-step league is already informative for that question, subject to its documented single-seed and optimizer-policy qualifications.

Strict matched ablation

This asks:

What changes when one mechanism axis is varied and the rest are fixed?

For the expansion question, I would hold constant:

• activation;
• read normalization;
• temporal/recency configuration;
• projection and RoPE layout;
• optimizer routing;
• learning-rate and update-scale policy;
• data order and training seed;
• validation and warmdown schedule.

Each arm would still need its whole-model parameter count recomputed through the benchmark adapter. I would additionally report:

• trainable mixer parameters;
• effective fast-state size;
• state bytes by dtype;
• estimated write/read FLOPs;
• measured training throughput.

This would prevent several different claims from collapsing into “same budget.” The public benchmark description is clear that the present league matches whole-model parameters by trimming the adjacent FFNs. That is a reasonable benchmark contract, but it is not the same as matching state size, compute, or memory traffic.

For the expansion arm, I would also distinguish two random variables:

• the normal training seed;
• the deterministic expansion_key.

A single training seed with a single fixed map cannot show whether a small result is robust to either source of variation.

The cleanest reporting pattern may therefore be:

These are complementary rather than competing evaluations.

I found the algebra easier to follow when “exact” is split into separate claims.

The ThetaScan algorithm notes already document the important GN qualifications:

• writes are linearized around a shared slow reference;
• the RMSNorm Jacobian rank-one correction is omitted;
• a diagonal/Jacobi-style Gram approximation is used;
• this is not full-Jacobian exact GN;
• summing token-local updates is not a joint nonlinear solve over all writes.

So a compact phrasing for future readers might be:

The defined token-local write statistics accumulate exactly under the model’s associative algebra; the GN write construction itself is a documented local approximation.

That also separates algebraic scanability from floating-point behavior. The JAX associative-scan documentation requires an associative combine operation, while PyTorch’s numerical-accuracy notes explain why mathematically equivalent reduction orders need not be bitwise identical.

There is also a useful boundary case in DeltaNet: a write can depend on the evolving state and still admit chunkwise parallelization when the transition has suitable linear structure. I would therefore describe ThetaScan’s distinctive choice more narrowly:

It keeps a nonlinear fast-weight read while defining token-local deposits at a shared slow reference, allowing the deposits to reduce through an ordinary additive scan.

That is more precise than treating state-independent writes as the only route to sequence parallelism.

For a future fused backend, a compact parity contract could include:

• sequential reference versus tree/chunk scan;
• materialized versus fused forward;
• backward gradients for inputs and parameters;
• FP32 and BF16;
• odd and non-power-of-two lengths;
• chunk-boundary cases;
• full sequence versus persistent streaming;
• duplicated samples alone versus inside a batch;
• very small denominator and inactive-feature cases;
• long positive accumulations and signed temporal mixtures.

The current portable-path parity results provide a useful baseline for that work.

The shared-reference construction buys order-independent deposits, but it raises a separate question: how well do independently computed local updates superpose after the nonlinear read?

In a tiny synthetic diagnostic, independently computed shared-reference writes were order invariant to machine precision. As the number and similarity of writes increased, however, their simultaneous fit degraded. Sequential relinearization could fit some of those cases better, but became order dependent.

I would not interpret that as a capacity result for the trained ThetaScan model. The diagnostic used a tiny random nonlinear memory rather than learned LM projections, and did not reproduce the complete normalized/recency architecture.

What it does provide is a useful test decomposition:

A. Taylor/Jacobian drift

Compare:

• the actual nonlinear output after applying the accumulated fast update;
• the first-order output predicted at the shared slow reference.

Sweep:

• number of writes;
• accumulated update norm;
• key overlap;
• threshold-boundary crossings;
• sequence position or recency branch.

Record:

• Taylor residual;
• Jacobian change;
• reconstruction error;
• update norm;
• active-feature changes.

Possible interpretations:

Small Taylor residual, poor retrieval:
    address geometry or collision is the stronger suspect.

Rapidly growing Taylor residual with update norm:
    the shared-reference linearization radius is the stronger suspect.

Abrupt changes near threshold crossings:
    the piecewise nonlinearity is interacting with the GN approximation.

B. Independent versus sequential updates

For a small controlled problem, compare:

1. independent writes computed at one shared reference;
2. sequentially relinearized writes;
3. the same sequential writes in reverse order;
4. optionally, a small joint optimization reference.

Record fit error, order gap, Jacobian drift, and update norm.

Sequential improves fit but has a large order gap:
    order invariance versus local refitting is the trade-off.

All methods fail only at high address overlap:
    representation/capacity is the stronger bottleneck.

Independent updates degrade much earlier:
    shared-reference superposition becomes the stronger candidate.

C. Separate context length from memory load

A single sequence-length sweep mixes several different failure modes. I would vary at least:

The first stage could operate directly on controlled key/query vectors, bypassing learned projections. The second stage could use learned representations and standard sequence tasks.

Zoology and MQAR are useful starting points for associative recall, while Just Read Twice is a reminder that presentation order itself changes the burden placed on a recurrent state. Exact-record recall should not be the only task, because ThetaScan’s normalized read is also naturally a query-dependent mixture. A multi-value test with a known expected mixture would exercise that part of the design more directly.

For normalized reads, I would log internal state as well as task accuracy:

• numerator norm;
• denominator quantiles;
• denominator-floor hits;
• active-feature fraction;
• target denominator share;
• effective support size;
• target-distractor margin;
• per-recency-branch contribution;
• head-wise effective rank.

A stable output norm alone would not establish low cross-talk: an increasing denominator may simply dilute all records together.

The current result seems most useful when read exactly as the repository describes it: an early, parameter-matched architecture-plus-policy comparison.

The league record and reproducibility notes already preserve several important boundaries:

• one training seed;
• attention and Mamba-3 use the host muon-2d policy;
• ThetaScan uses muon-2d+theta;
• the completed result is a protocol reproduction target rather than a bitwise replay guarantee;
• the historical research map and a map regenerated from the public expansion_key need not be bitwise identical;
• the reference benchmark uses the portable quadratic evaluator;
• optimized kernels and streaming inference remain early work.

For optimizer controls, I would compare routing rather than only optimizer names:

• which matrix-shaped memory factors receive Muon;
• which vectors/scalars receive the auxiliary optimizer;
• per-group learning rates;
• weight decay;
• update scaling;
• treatment of threshold parameters.

A full optimizer grid would probably be excessive before the activation/expansion contrast is resolved. One representative crossed control among attention, dense GN, and expanded GN may be enough to show whether the current ranking is sensitive to the policy assignment.

There is also an API boundary relevant to long-context reproduction: the current public forward path does not yet expose persistent fast state across separate calls. The roadmap identifies persistent streaming state as future work. Until that exists, evaluating one complete sequence in one call is not equivalent to splitting it across independent forward calls.

I do not think ThetaScan needs to be forced into one existing category, but a few neighboring lines of work provide useful axes for comparison.

• Test-Time Regression separates associative-memory designs into the regressor class, association weighting, and test-time optimizer. In that vocabulary, ThetaScan-GN can be viewed as a nonlinear regressor receiving token-local approximate updates at a shared reference.

• DeltaNet is a useful contrast for state-dependent corrective writes and structured sequence parallelism.

• MesaNet provides a different point in the design space: it spends additional computation to approach a locally optimal cumulative regression solution rather than adding independently constructed local updates.

• Zoology/MQAR provides a controlled route into associative-recall evaluation, while also motivating movement beyond the simplest synthetic lookup tasks.

I would use those as comparison vocabulary, not as claims of equivalence or priority.