Experimenting with hypersurface-constrained dynamic weight updating [P]

Hey everyone!

I want to share the results of a small language model architecture experiment I’ve been running as a side project. The ultimate goal is to test an approach to reduce the number of the model's training parameters, since the main bottleneck in training is VRAM.

The core idea is very similar to Universal Transformer. Let's take just a single decoder block and iteratively pass the input through it L times in a loop. But instead of updating the input with the loop iteration depth information, my model dynamically updates the weights of the base layer.

How is it implemented?

The model uses a set of learned hypersurfaces to generate weight deltas (𝛥Wl) from the hypersurfaces' cross-sections. So a given weight matrix at layer l is constructed as Wl = W0 + 𝛥Wl.

The hypersurfaces are defined by a series of periodic functions. The best results I've got so far are from using a triangular wave. The model learns the amplitudes, frequencies, and phases of these waves across the coordinate dimensions. Given the size of this set E this gives a total of 3*E*dim parameters.

Initially, I tried generating the full weights purely from hypersurfaces, but it was too restrictive and failed to converge. Thus, I moved to my current approach where the model has a base decoder layer that is being updated.

I also added a state vector that is calculated from the input sequence using Gated Linear Attention. It is used to modulate the hypersurface geometry during the forward pass, making the generated weight deltas 𝛥Wl sequence-aware.

Pre-training results

I ran the pre-training experiments on a 10B-token sample from the FineWeb-Edu dataset. I wanted to keep the training parameter count minimal, so I used a pre-trained, frozen embedding layer from GPT-2. I also do not use any positional encoding, following the NoPE approach. The sequence length was 1024 tokens, the batch size was 16, and I ran the training for 10,000 steps.

I used 3 baselines:

• A simple 1-layer decoder-only transformer
• A single-layer transformer unrolled across 24 loop iterations
• A standard 24-layer decoder-only transformer

The experimental models:

• A single loop block using sinusoidal surface deltas, no context modulation
• A single loop block using triangular wave surface deltas with context modulation
• A model that uses three stacked loop blocks, each utilizing triangular wave surface deltas and context modulation

Here's a training loss chart:

Training results

Model parameter count

A standard 24-layer decoder-only transformer: 169,906,944

3 loop blocks + triangular wave + context modulation: 27,162,624 (~16% of the baseline model's size)

Key Takeaways:

While the classic decoder-only architecture still produces the best absolute loss, the Triangular Surface + Context model shows a real performance boost over a standard unrolled baseline.

This model introduces a whole new set of parameters and hyperparameters. I am currently experimenting with the size of the triangular wave function set and the size of the context vector. I am also trying different initialization strategies for the hypersurface parameters. The (very vague) goal isn't necessarily to beat a full-parameter model, but to find a configuration that is "good enough" in terms of loss while remaining drastically lighter on hardware resources.

Another possibility is that hypersurface representations have different convergence dynamics and require more steps and more tokens to converge. So my next step is to train a model with 3 loop blocks and a larger context vector, and an expanded functional basis on the whole 10B-token sample from the FineWeb-Edu dataset.

GitHub repo: https://github.com/morgengramlich/LoopSLM