Two AI Architecture Experiments: Shared MLP and FFT-Based Nonlinear Memory

A developer shared two small neural network architecture experiments on Reddit—PHSMLP and Pilot. Both are single-seed, small-scale research explorations, and the author explicitly states they don't claim these are broadly better than Transformers, but some of the results are worth noting.

PHSMLP: Making Transformers Reuse the Same MLP

The core question behind PHSMLP (Pooled Head MLP) is: can a Transformer share a single MLP across all layers instead of having one per layer? If sharing reduces parameters, can learned head pooling make up for the performance loss?

The author trained four models on 500 million FineWeb-Edu tokens, with the following results (lower perplexity is better):

Model	Perplexity	Accuracy
Standard Transformer	88.75	23.95%
Standard + pooling	90.60	23.50%
Shared MLP	105.98	21.18%
Shared MLP + pooling	105.66	21.03%

The shared MLP trails the standard Transformer by less than 3 percentage points in test accuracy while reducing parameters by 13.4%. That gap isn't large for an aggressive parameter-sharing experiment.

Pooling shows localized value within the shared architecture

Although pooling didn't let the shared model catch up to the standard Transformer, it showed some positive effects within the shared architecture. Compared to the pure shared MLP, adding pooling resulted in:

• Test perplexity dropping from 105.98 to 105.66
• Fact overwriting accuracy rising from 29.7% to 32.8%
• Multi-hop reasoning accuracy rising from 28.1% to 30.7%

More convincing evidence comes from an ablation experiment: disabling pooling in the trained shared model caused perplexity on the ablation test subset to spike from 100.02 to 181.20. This shows the model genuinely learned to use the pooling mechanism meaningfully rather than treating it as a redundant component.

The author also acknowledges that compared to an independently trained shared MLP, the gains from pooling are small and not entirely conclusive. But the combination of a shared MLP retaining most of the original Transformer's performance, plus pooling providing a useful learning pathway, is worth further investigation.

Parameter Sharing is a common technique in neural networks for reducing model parameters and improving generalization. Its core idea is to let multiple positions or modules in the network reuse the same set of weights. In Transformers, typical parameter sharing includes sharing self-attention weights across layers, sharing feed-forward networks (FFN/MLP), and so on. PHSMLP explores exactly this kind of cross-layer MLP sharing: a standard Transformer has an independent MLP in each layer, with parameter count growing linearly with depth; if all layers share the same MLP, this portion of parameters can be compressed to a constant. The cost is limited model expressiveness, since different layers can't learn their own specialized feature transformations. The degree of performance loss from shared MLPs depends on task complexity and model depth—smaller losses in shallow or simple tasks, more noticeable losses in deep tasks or those requiring hierarchical abstraction.

Head Pooling in the PHSMLP context refers to a mechanism for weighted aggregation of multiple attention heads or inter-layer representations. In multi-head attention, each head attends to a different subspace; head pooling combines the outputs of these heads through learned weights, allowing the shared MLP to receive filtered, more informative input. Unlike simple average pooling or max pooling, learnable head pooling lets the model adaptively decide which heads' information is more important, thereby retaining some expressiveness under the constraints of parameter sharing.

Pilot: FFT-Accelerated Explicit Input History + Nonlinear Matrix Memory

Pilot builds on the author's earlier ILRM research, exploring whether FFT can be used to compute explicit input history without constructing a temporal Toeplitz matrix, while retaining an independent nonlinear recurrent memory.

The architecture has three components:

• Exponential history of prior input matrices
• Nonlinear matrix memory controlled by access and write gates
• Key, value, and query projections for constructing input and reading memory

History computation can be done with FFT, sequential scan, or parallel prefix scan—all three compute the same mathematical history, and numerical checks show agreement within floating-point tolerance. The key distinction is that FFT operates on the history pathway, not the nonlinear recurrent matrix itself. The memory doesn't need to be Toeplitz, diagonal, or linear.

Wins on language modeling, loses on memory tasks

Using two T4 GPUs and 63.57 million FineWeb-Edu tokens per model, the author compared Pilot FFT against a parameter-matched four-head Transformer:

Model	Parameters	Perplexity	Accuracy
Transformer	428,608	126.24	17.61%
Pilot	428,523	105.15	20.35%

Pilot's test perplexity is 16.7% lower than the Transformer's, which the author calls a surprise. But it doesn't win across the board on memory and reasoning tests: of 96 multiple-choice questions, the Transformer answered 17 correctly while Pilot got only 11. Both models scored zero on unrestricted next-token answers and neither could reliably solve recall, fact binding, or multi-hop reasoning tasks.

Pilot also makes limited use of long-range context and trains about 60 times slower in the current implementation. The positive language modeling results coexist with shortcomings in long-term memory and training efficiency.

As for the name, the author says the AI field already has plenty of fierce animal names and wanted something with a British connection—Britain doesn't have many formidable beasts to choose from, so the pilot whale came to mind, and it also happens to be a pilot experiment for the idea itself.

FFT (Fast Fourier Transform) is used in sequence modeling to accelerate convolution operations. Convolution in the time domain is equivalent to pointwise multiplication in the frequency domain, so applying FFT to a sequence, multiplying by the frequency-domain representation of a filter, and then doing an inverse FFT reduces convolution complexity from O(N²) to O(N log N). Pilot leverages this property to compute "explicit input history": history information is essentially a convolution of the input sequence with some kernel function, and constructing a Toeplitz matrix directly in the time domain requires O(N²) storage and computation, whereas the FFT path avoids explicit matrix construction. Note that FFT can only accelerate linear convolution; the nonlinear recurrent memory portion still requires sequential or parallel scan processing.

Toeplitz Matrix is a matrix where each diagonal has the same elements, often appearing in representations of convolution and linear time-invariant systems. In sequence models, if causal convolution is represented as matrix multiplication (current output is a weighted sum of historical inputs), that matrix is a lower-triangular Toeplitz matrix. Constructing and storing a full Toeplitz matrix is expensive, so Pilot chooses to compute the equivalent result implicitly via FFT.

Perplexity is the most commonly used evaluation metric in language modeling, defined as the exponential of the model's average negative log-likelihood for each token in the test set. Lower perplexity means the model predicts the next token more accurately and is less "perplexed." Its value can be understood as how many candidate words the model needs to choose from on average at each step—ideally equal to the vocabulary size, but in practice far below it. Perplexity is sensitive to tokenization and vocabulary size, so cross-model comparisons require ensuring the same tokenizer and test set are used. Pilot's perplexity of 105.15 versus the Transformer's 126.24 indicates that at the same parameter count, Pilot fits the statistical patterns of language sequences better.

Multi-hop Reasoning refers to questions that require chaining together multiple facts or reasoning steps to answer, such as "What is the capital of the country where someone was born?" These tasks require the model to locate and combine multiple scattered pieces of information in context, placing high demands on long-range dependencies and memory mechanisms. Both Pilot and the Transformer perform poorly on these tasks, indicating that current architectures still have clear shortcomings in explicit memory and compositional reasoning.

The Common Thread Between the Two Experiments

The author repeatedly emphasizes that these are single-seed, small-scale experiments requiring more testing. PHSMLP demonstrates that a shared MLP can retain most of a Transformer's prediction performance with fewer parameters, and pooling shows encouraging localized gains within the shared architecture; Pilot demonstrates that FFT-compatible input history can work synergistically with nonlinear recurrent matrix memory, delivering language modeling improvements at matched parameter counts.

Code, research PDFs, charts, and results for both works are publicly available in a GitHub repository. The author is particularly interested in learning how to improve Pilot's memory and speed without adding substantial extra mechanisms or changing the original mathematics. For researchers interested in efficient architectures and alternative attention mechanisms, these two directions offer reproducible starting points.

One-line Summary

Two small architecture experiments: PHSMLP with a shared MLP and Pilot with FFT-accelerated memory, each with localized highlights but neither broadly surpassing Transformers.

Paragraph Summary

The developer shared two single-seed, small-scale neural network architecture experiments, PHSMLP and Pilot. PHSMLP explores sharing a single MLP across layers, trailing the standard Transformer by less than 3 percentage points in test accuracy after training on 500 million tokens while reducing parameters by 13.4%, with ablation experiments proving the pooling mechanism is effectively utilized. Pilot combines FFT-computed explicit input history with nonlinear matrix memory, achieving 16.7% lower language modeling perplexity than a parameter-matched four-head Transformer, but performing worse on multiple-choice memory tasks, training about 60 times slower, and making limited use of long-range context. The author emphasizes both works are preliminary explorations, with code and PDFs publicly available, and hopes the community can help improve Pilot's memory and efficiency.