Hi. I think this has become much easier to experiment with. I played with it a little:
The biggest thing for me is that turning WND into WiND makes several of the questions from the earlier architecture thread much easier to test directly.
After going through the current code and doing a few small CPU-level checks, I think I would start with two fairly practical things before changing the architecture further:
- make one tiny dense CPU path the known-good entry point: install → build → one forward/backward;
- make the execution contract of the generic
wind.build(...) path versus LanguageModel(...) explicit.
Those two paths currently differ in some fairly important ways. That is not necessarily a problem — they may intentionally serve different experiments — but documenting the distinction would make external experiments much easier to interpret.
For reproducibility, the checks below were against WiND commit 6df9170.
1. The generic API and LanguageModel currently describe rather different experiments
The most important difference I found is the meaning of depth × iterations.
In the generic ReasoningDepth, one iteration selects one layer:
for index in range(self.iterations):
layer = self.layers[index % len(self.layers)]
...
x = layer(x)
So, for example:
depth = 4
iterations = 1 -> layer 0
iterations = 4 -> layers 0,1,2,3
iterations = 8 -> layers 0,1,2,3,0,1,2,3
I also checked this with backward hooks: with four Depth layers and iterations=1, only the first layer executed and only that layer received gradients.
LanguageModel has different semantics: each reasoning iteration runs the full Depth stack.
Conceptually that is closer to:
depth = 4
iterations = 1 -> 4 layer applications
iterations = 4 -> 16 layer applications
This matters especially because the generic Config defaults are currently depth=6, iterations=1. Someone reading that as “a six-layer Depth stack, repeated once” would be running a different experiment from what the current generic implementation does.
If this distinction is intentional, I think even a short note in the API docs would remove a lot of ambiguity.
There is another related difference around the state/Bank boundary.
In the generic builder, the default Compressor and Feature Bank both use the bank_tokens budget. In other words, the mutable reasoning stream and the reference Bank do not have independently configurable token budgets by default.
LanguageModel, on the other hand, exposes separate:
state_tokens
bank_tokens
which seems closer to the original idea of having a smaller iterative state reasoning against a richer reference representation.
I would therefore record at least these fields alongside any result:
frontend:
depth:
iterations:
executed layer applications:
state tokens:
bank tokens:
bank construction:
bank retrieval rule:
bank differentiable?:
bank mutable during recurrence?:
That would make comparisons across WiND experiments much easier for future readers.
One other terminology detail that may be worth making explicit: “read-only” can mean several different things.
A Bank can be:
- immutable during the recurrent computation;
- request-local rather than persistent across requests;
- detached from autograd.
Those are separate properties.
The generic FeatureBank defaults to a detached, request-local, non-written Bank. The LanguageModel learned-query Bank is also non-written during the Depth loop, but remains differentiable end-to-end.
So I would probably avoid treating “read-only” and “stop-gradient” as synonyms in experiments.
2. One small long-context distinction in the generic Feature Bank
There is also a potentially useful control for longer inputs.
The generic FeatureBank currently keeps:
features[:, :max_tokens]
when the encoded sequence is longer than the Bank budget.
The Compressor behaves differently: it uses adaptive pooling over the whole encoded sequence.
So for a long input, it is possible to have:
initial reasoning state
<- information pooled from the whole input
Feature Bank
<- only the first N encoded positions
That means a simple “Bank enabled vs Bank disabled” experiment on long inputs could mix together two effects:
- whether access to an immutable reference is useful;
- whether the relevant information happened to be covered by the Bank selection rule.
A very cheap sanity check would be to move the same informative content between the beginning and end of the input and compare the Bank intervention effect.
If that positional sensitivity appears, then it becomes worth separating “reference memory capacity” from “reference coverage policy.”
LanguageModel does not have exactly the same issue because its Bank is produced with learned-query attention over the source rather than prefix truncation.
3. A tiny clean-environment path would probably help people try WiND
Since this post is explicitly inviting people to experiment with and modify WiND, I think a tiny known-good CPU example could have unusually high value.
Something approximately like:
from wind import Config, build
import torch
model = build(
Config(
dim=64,
heads=4,
depth=2,
iterations=2,
bank_tokens=8,
wide_type="dense",
)
)
x = torch.randn(2, 16, 64)
y = model(x)
y.square().mean().backward()
print(y.shape)
with one tested install command would give people a clean baseline before they enter the PKM/CUDA path.
I mention this because I hit a packaging boundary while checking the current repository.
In a clean wheel-build smoke test:
- the root
wind wheel did not contain flashpkm;
- the nested FlashPKM project built a metadata-only wheel in that layout;
- consequently the optional PKM import was unavailable in that clean test.
I would not read too much into this as an architectural issue — it looks more like a packaging/discovery boundary — but it makes an explicit dense-vs-PKM installation split useful.
The current bridge already has the right conceptual split:
dense CPU path
-> no FlashPKM required
PKM path
-> FlashPKM / Triton / CUDA setup
so documenting those as two separate quick starts may be enough.
For reference, setuptools’ package-discovery rules match include patterns against the full import-package name, which is relevant when a package is mapped from a nonstandard directory layout:
setuptools package discovery
I would probably also make the first example self-contained rather than dependent on a local dataset path. Once the one-batch CPU path works, the larger training example becomes much easier to debug.
4. The next architecture test I would try is a same-checkpoint causal test
Going back to the earlier WND discussion, I think the most informative next step is no longer another architectural variant. It is checking what an already-trained model actually depends on.
For one fixed checkpoint, I would compare something like:
A. normal Bank read on every iteration
B. Bank read only on the first iteration
C. no Bank read
D. shuffled / swapped Bank
E. normal Bank, but iterations = 1 / 2 / 4 / 8 / ...
Ideally split by problem difficulty rather than only reporting one aggregate score.
That gives a fairly useful decision tree:
More iterations help, especially on harder examples
-> evidence that recurrent compute is doing useful conditional work
Removing Bank reads hurts
-> evidence that the iterative state depends on the input-conditioned reference
Iterations help, but Bank removal barely matters
-> the recurrence may mostly be operating inside the compressed state
Bank access matters, but more iterations do not
-> reference retrieval may be useful without much iterative refinement
Neither intervention matters much
-> the task may already be solved upstream,
or the current training setup may not be forcing the Depth path to matter
This seems especially relevant after the earlier “Encoder does all the work” / shortcut discussion.
I would be cautious about calling either part “knowledge” or “reasoning” based only on where it sits in the architecture. The stronger evidence would be behavioral:
- what information is available in each representation;
- which intervention changes performance;
- whether sensitivity changes with task difficulty;
- how the latent trajectory changes across iterations.
That turns “Wide = memory, Depth = reasoner” from a naming assumption into something experimentally testable.
5. There is some useful nearby work for designing those controls
A few recent directions seem especially relevant, although none of them is exactly WND.
Looping versus memory capacity
Adaptive Loops and Memory in Transformers: Think Harder or Know More? explicitly separates iterative computation from additional learned memory capacity.
One result I found useful conceptually is that the two resources do not behave like simple substitutes: looping helped more on mathematical reasoning, while memory helped recover performance on commonsense-style tasks.
The memory mechanism there is not the same as the WND Feature Bank, so I would not transfer the result directly. But the experimental decomposition seems very useful:
more recurrent compute
!=
more storage/reference capacity
That suggests varying iterations, state_tokens, bank_tokens, and Wide/PKM capacity independently when possible.
Re-reading the input during recurrence
Logical Extrapolation Without Overthinking found that recurrent models could degrade when iterated far beyond their training regime, and used a recall mechanism that repeatedly makes the original problem instance available.
WND’s Bank is not the same mechanism, but it suggests another interpretation of repeated Bank access:
it may be useful not only as “knowledge lookup,” but also as an immutable input-derived reference that prevents the recurrent state from having to preserve every relevant detail internally.
That is why I think:
every-iteration Bank read
first-iteration-only Bank read
no Bank read
could be more informative than only comparing Bank sizes.
Does architectural separation actually produce functional specialization?
One Model, Two Roles: Emergent Specialization in a Shared Recurrent Transformer / code is also interesting here.
AIR deliberately studies whether recurrent states develop different functional roles, and uses interventions such as input-injection asymmetry, state freezing, decoded trajectories, and attention analysis rather than relying only on module names.
Again, its architecture is different from WND, but the methodology maps nicely onto the original question:
structural separation
-> hypothesis
intervention / freeze / trajectory sensitivity
-> evidence for functional separation
I think that distinction would be valuable as WND experiments become larger.
6. For comparisons, I would report executed compute separately from parameter count
Because recurrence reuses weights, parameter count alone can make comparisons difficult to interpret.
I would separate at least:
| Quantity |
What it tells you |
| Unique parameters |
storage / parameter efficiency |
| Executed layer applications |
actual recurrent depth |
| FLOPs |
compute budget |
| Wall-clock latency |
sequential cost in practice |
| Peak VRAM |
deployment/training cost |
| State tokens |
mutable workspace capacity |
| Bank tokens |
reference capacity |
| KV/cache footprint if relevant |
inference memory cost |
| Seeds |
stability |
A useful recent example is SMELT: Scaling Laws for Compute-Matched MoE Looped Transformers, which explicitly matches per-token FLOPs, non-embedding parameter count, and KV cache when comparing looped and non-looped Transformers.
I do not think WND needs that level of benchmark machinery yet. But keeping these axes separate early would prevent later results from becoming hard to interpret.
For example:
same parameters + more iterations
tests something different from:
same FLOPs + fewer unique parameters
and both are different again from:
same latency / VRAM envelope
The earlier WND discussion was already moving in this direction, so WiND now gives a convenient place to make those accounting rules explicit.
7. PKM looks like a good place for utilization diagnostics, not just capacity numbers
If you continue with PKM as the Wide side, I think it would be useful to distinguish nominal memory capacity from actually used capacity.
The current FlashPKM path exposes enough retrieval-side auxiliary information that you could inspect things like:
- fraction of slots ever selected;
- cold-slot fraction;
- retrieval-frequency concentration;
- entropy of the selected distribution;
- differences across heads;
- how these change with task type or training stage.
I would not infer much from those numbers at random initialization.
But on a trained model they could answer a useful question:
Is increasing PKM memory_size giving the model a larger usable memory,
or mostly a larger address space that routing does not actually occupy?
That is probably more informative than reporting memory size alone.
8. Two small implementation notes I would keep separate from the main architecture discussion
These are much narrower than the points above, so I would treat them as implementation checks rather than architecture criticism.
AdaptiveFeatureBank scorer
In a minimal backward check with detach=False, I got no task-loss gradient on:
AdaptiveFeatureBank.score.weight
The reason seems to be that the scorer output is currently used to choose integer topk(...).indices, after which the selected feature values are gathered. The scorer value itself does not remain in the continuous downstream computation.
This is different from the generic Retrieval path: there, top-k chooses indices too, but the selected similarity scores are then used in a softmax weighting, so gradients still flow through the selected scores and q/k/v projections.
So if AdaptiveFeatureBank.score is intended to learn directly from the normal task loss, I would double-check that contract.
If it is intended to be externally supervised, fixed, or trained with an auxiliary objective, then this may simply need documentation.
alpha learning
I also did a very small random-initialization sanity check on use_alpha_learning.
At initialization its output effect and alpha gradient were extremely small.
That makes some sense from the current computation:
state
vs
RMSNorm(state)
blend them with alpha
then apply RMSNorm again
With the initial normalization parameters, those two vectors are already very closely related, and the final normalization removes much of the remaining scale difference.
This does not show that alpha stays ineffective after training — learned non-uniform norm weights could change the picture substantially — so I would not draw a conclusion from the initialization test.
It is just something I would monitor if alpha is used in a quality experiment.
Overall, I think WiND changes the situation quite a bit compared with the earlier thread: the interesting questions are now much easier to turn into small controlled experiments.
My default route would be:
1. make a tiny dense CPU path reliably reproducible
2. document generic-vs-LanguageModel execution semantics
3. choose one frontend and record its state/Bank/iteration contract
4. use same-checkpoint Bank/iteration interventions
5. only then decide whether the next bottleneck is
recurrence, state capacity, Bank coverage, retrieval, or Wide/PKM capacity
That seems like a fairly low-cost way to preserve the original WND idea while making it progressively easier to tell which part is actually doing what.