Pre-Training a Hybrid U-Net Transformer

During my batch at the Recurse Center, I’ve been experimenting with language model pre-training on a limited compute budget. I wanted to try changes to the Transformer architecture while also working on the training code, since faster runs would give me more room to compare ideas.

I spent most of that time on skip connections inspired by U-Net, using them to combine representations from earlier and later Transformer layers.

Why U-Net?

U-Net was introduced for biomedical image segmentation. Its encoder reduces spatial resolution as it builds features, and its decoder combines those features with earlier ones through skip connections. That lets the output use both local detail and information gathered at lower resolutions.

I wanted to test a similar connection pattern in a language model. My version keeps the sequence length fixed and pairs layers by depth, so it borrows U-Net’s skip structure without reproducing its spatial downsampling and upsampling.

I split 16 Transformer layers into two groups. Layers 1 through 8 save their activations, and layers 9 through 16 combine them with their own activations through learned gates.

Each gate has one learned parameter per hidden dimension, initialized to 0.1. It combines the two inputs as follows:

\[\text{output} = \text{gate} \times \text{skip} + (1 - \text{gate}) \times \text{current}\]

A larger gate gives the saved activation more weight, while a smaller gate gives the current activation more weight. Because each layer has its own gate, training can adjust these contributions separately.

The mixing module is small:

1
2
3
4
5
6
7

class GatedMixer(nn.Module):
    def __init__(self, d_model: int) -> None:
        super().__init__()
        self.gate = nn.Parameter(torch.full((1, 1, d_model), 0.1))

    def forward(self, skip: torch.Tensor, current: torch.Tensor) -> torch.Tensor:
        return self.gate * skip + (1 - self.gate) * current

The saved activations are used in reverse order: layer 9 receives layer 8’s output, layer 10 receives layer 7’s, and so on. This creates connections between matching positions in the two halves of the model.

Not all layers are equal

I also split the attention mechanisms between the two halves. The idea was to let the earlier layers work with nearby tokens, then give the later layers access to the full preceding context.

Layers 1 through 8 use sliding-window multi-head attention, following the local-attention idea used in Longformer. Each query is limited to a window ending at its own position, which preserves the causal constraint.

In the window described here, position \(i\) can attend to positions in \([i - 256, i]\). These layers don’t directly attend to distant tokens; the upper layers provide the wider context.

Layers 9 through 16 use multi-query attention (MQA). The query heads share key and value projections, reducing the amount of key and value data that needs to be stored during autoregressive inference.

With 8 heads and d_model=1024, each head has 128 dimensions. Standard multi-head attention stores a separate key and value for each head, while MQA shares one pair across all eight. At the same sequence length, the MQA layer therefore uses about one eighth as much KV cache storage.

Both attention types use RoPE to represent position. RoPE rotates query and key vectors according to their positions, so their dot product also depends on the relative distance between them.

The rotation is computed using sine and cosine functions:

\[\text{freq}_{j} = \frac{1}{10000^{2j/d}}\]

\[\text{RoPE}(x, \text{pos}) = \begin{pmatrix} x_{1} \\ x_{2} \\ \vdots \end{pmatrix} \odot \begin{pmatrix} \cos(\text{pos} \cdot \text{freq}) \\ \cos(\text{pos} \cdot \text{freq}) \\ \vdots \end{pmatrix} + \text{rotate\_half}(x) \odot \begin{pmatrix} \sin(\text{pos} \cdot \text{freq}) \\ \sin(\text{pos} \cdot \text{freq}) \\ \vdots \end{pmatrix}\]

The rotation can be calculated beyond the sequence lengths used in training, but that doesn’t guarantee the model will work well at those lengths. I would still need to evaluate longer-context inference separately.

Letting activation functions learn

The feedforward blocks start with squared ReLU, \(\text{ReLU}(x)^2\). The Primer paper includes this change among the modifications that improved its language models. I wanted to see what would happen if the exponent could vary by layer.

I initialized a learned exponent at 2 and allowed training to adjust it:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

class SquaredReLUFFN(nn.Module):
    def __init__(self, d_model: int, expansion: float) -> None:
        super().__init__()
        hidden = int(d_model * expansion)
        self.w1 = nn.Linear(d_model, hidden)
        self.w2 = nn.Linear(hidden, d_model)
        self.alpha = nn.Parameter(torch.tensor(2.0))  # learnable exponent

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.w1(x)
        x = torch.relu(x)
        alpha = torch.clamp(self.alpha, 0.5, 4.0)  # keep it reasonable
        x = torch.pow(x + 1e-6, alpha)
        x = self.w2(x)
        return x

Some layers settled near \(\alpha \approx 2.3\), while others stayed around 2.0 or moved toward 1.8. That showed the parameter was changing during training, although comparing against a fixed-exponent model is still necessary to tell whether it helps the final result.

I also varied the FFN expansion ratio by layer:

• Layers 1-8: 4× expansion (4096 hidden for d_model=1024).
• Layers 9-16: 2.5× expansion (2560 hidden).

This gives the lower half more feedforward capacity while keeping the upper half smaller. It was a design choice to test alongside the attention split, rather than something I knew would be better in advance.

Making it actually run

A 16-layer model with a hidden size of 1024 was large enough that training time and memory use constrained the experiments. I worked through several optimizations to make repeated runs practical.

I started with bfloat16 mixed precision. Autocast runs eligible operations at lower precision while other operations and the parameter gradients can remain in fp32:

1
2
3
4

with torch.amp.autocast(device_type='cuda', dtype=torch.bfloat16):
    logits = model(input_ids)
    loss = cross_entropy(logits, targets)
loss.backward()

Bfloat16 has the same 8-bit exponent range as fp32, with 7 bits of fraction precision instead of 23. That wider range was useful for avoiding some overflow problems associated with fp16. It still loses precision, so the loss and gradients needed checking during training.

I also used NVIDIA’s Transformer Engine to run supported attention operations in FP8. A value stored in FP8 takes half the space of bf16, although the memory savings for the whole model depend on which tensors use it. In these experiments, the FP8 setup gave roughly 30 to 40% more throughput and allowed larger batches:

1
2
3
4

if self.use_fp8 and te is not None:
    hidden_states = x.transpose(0, 1).contiguous()
    with te.fp8_autocast(enabled=True):
        output = self._te_mha(hidden_states)

Those results came from an H100. The available FP8 kernels depend on the GPU and the Transformer Engine version, so the same setup needs checking before moving it to other hardware.

FlashAttention reduces attention’s memory traffic by calculating it in blocks without storing the full attention matrix. PyTorch’s scaled_dot_product_attention can select a fused implementation when the inputs and hardware support it:

For the batch sizes I tested, attention memory use fell by 5 to 8× and end-to-end training was roughly 2 to 3× faster.

The sliding-window mask in my implementation didn’t use the same fast path, so layers 1 through 8 fell back to standard attention. That limited the speedup, even with the smaller attention window.

I also enabled these backend settings:

1
2

torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.benchmark = True

TF32 (TensorFloat-32) lets supported matrix multiplications use reduced mantissa precision while retaining fp32’s exponent range. I measured a 15 to 20% speedup in my H100 tests. The setting affects eligible fp32 operations, so its contribution depends on how much of the workload already runs in bf16 or FP8.

cuDNN benchmarking searches for a suitable convolution algorithm for the input shapes. It only matters where the model uses convolution operations; it doesn’t accelerate the attention layers themselves.

torch.compile captures parts of the model for compilation and can fuse operations to reduce overhead:

1

model = torch.compile(model, mode='reduce-overhead', dynamic=True)

The reduce-overhead mode uses CUDA graphs where applicable to reduce Python overhead, while dynamic=True asks the compiler to handle changing shapes. Compilation made the first iteration slow, and some operations needed extra work to compile, but steady-state throughput improved by about 20 to 30%.

Gradient checkpointing reduced the activations I needed to keep in memory. It recomputes them during the backward pass:

1
2
3
4

if self.gradient_checkpointing and (i % self.checkpoint_every == 0):
    h = checkpoint(lambda inp: self.blocks[i](inp, attn_mask), h, use_reentrant=False)
else:
    h = self.blocks[i](h, attn_mask)

Checkpointing every layer reduced activation memory by roughly 70%. Iterations took about 20% longer, but the extra memory allowed larger batches, which made the trade worthwhile for this setup.

Optimizer choices

I tried Muon in modded-nanogpt, which applies an orthogonalization step to momentum-based updates for matrix weights. The optimizer’s description explains that step and which parameters should use a conventional optimizer. I used an AdamW group for the parameters outside the Muon group:

1
2
3
4

optimizer = MuonWithAuxAdam([
    {"params": muon_params, "lr": 8e-4, "use_muon": True},
    {"params": adamw_params, "lr": 8e-4, "betas": (0.9, 0.95)}
])

I used different peak learning rates for different parts of the model:

• Embeddings: 2e-3.
• Layers 1-8: 8e-4.
• Layers 9-16: 6e-4.
• LM head: 4e-4.

These settings came from trial and error. Higher rates worked better for the embeddings in these runs, while the upper layers and output head were more stable with smaller updates.

The schedule used 1,000 warmup steps followed by cosine decay to 10% of each group’s peak learning rate:

During warmup (\(t < t_{\text{warmup}}\)):

\[\eta(t) = \eta_{\max} \cdot \frac{t}{t_{\text{warmup}}}\]

After warmup (\(t \geq t_{\text{warmup}}\)):

\[\eta(t) = \eta_{\min} + \frac{1}{2}(\eta_{\max} - \eta_{\min}) \left(1 + \cos\left(\pi \frac{t - t_{\text{warmup}}}{t_{\text{total}} - t_{\text{warmup}}}\right)\right)\]

Warmup increases the step size gradually at the start of training. Cosine decay then reduces it as training progresses, allowing smaller updates toward the end.

To increase the effective batch size without keeping all its activations in memory at once, I accumulated gradients over two micro-batches before updating the weights:

1
2
3
4
5

for micro_step in range(2):
    loss = forward_and_loss(batch) / 2  # scale the loss
    loss.backward()
optimizer.step()
optimizer.zero_grad()

Dividing each loss by the number of accumulation steps keeps the accumulated gradient on the intended scale.

How training went

I trained on Wikipedia and OpenWebText, using about 10B tokens over 50,000 steps.

About the plots: these charts are illustrations generated by the page, rather than exported training logs. The learning-rate chart shows the configured schedule; the loss, throughput, gate, and memory charts use generated values.

This illustration shows training and validation loss decreasing together, with a small gap between the two curves.

The learning-rate plot follows the schedule in the code: 1,000 steps of warmup, followed by cosine decay. The rates for the different parameter groups remain proportional throughout.

The throughput illustration shows an initial increase followed by a steady rate of around 12k tokens per second.

The next illustration shows how different layers could settle on different gate values:

In this illustration, layers 9 and 10 have gates around 0.5 to 0.6, while layers 14 through 16 have values around 0.1 to 0.25. Higher gate values give more weight to the saved activations.

The memory illustration separates allocated from reserved memory and marks the H100’s 80GB capacity. The batch size, precision, and checkpointing settings all affect how much memory a run needs.

What actually helped

In my comparisons, the changes had the following effects. These results are specific to this model and training setup.

The U-Net skip connections improved loss by roughly 5% in my comparisons. The gate values differed across layers, which was useful to inspect, but their variation alone doesn’t explain the loss improvement.

Combining MHA in the lower half with MQA in the upper half changed training speed little. Inference was about 15% faster, with the smaller KV cache in the MQA layers helping reduce the amount of data involved.

Using different learning rates by layer gave an improvement of around 8% in the comparison, more than the 2 to 3% I had expected. The higher embedding rate accounted for an important part of that result.

The FP8 attention setup increased throughput by about 35%, and the reduced memory use allowed me to double the batch size.

After compilation, torch.compile gave roughly a 25% speedup. That was useful over a long run, despite the initial compilation cost.

Checkpointing’s 70% reduction in activation memory was also useful despite the slower iterations. It let me use a reasonable batch size without exceeding the GPU’s memory.

With these changes, I could train the 300M-parameter model on a single H100 over a few days, which was enough to make further experiments practical.

Things that flopped

Several other changes added complexity without improving the result.

Auxiliary loss at layer 8. I added a second language-model head halfway through the network and trained it on the same targets. Its learned loss weight shrank toward zero, so it contributed little to the final training objective.

Bigger sliding windows. Increasing the window to 512 or 1024 made training slower without improving loss in these runs, so I kept the 256-token setting.

Gate initialization. I tried starting the gates at 0.5 and 0.0 as well as 0.1. The results were similar enough that I kept 0.1.

Stochastic depth. Randomly dropping layers, following the stochastic depth idea, hurt convergence in this 16-layer model.

Code

The implementation is in my llm repository. The main files are:

• model/transformer.py: Hybrid U-Net Transformer architecture.
• model/attention.py: MHA and MQA with FP8 support.
• model/ffn.py: Squared ReLU with learnable exponent.
• train/loop.py: Training loop with all the optimizations.
• train/optim.py: Muon optimizer and learning rate schedules.

This is experimental code. It still lacks distributed training and a complete checkpointing system, and some parts need more testing before I’d use them for a long unattended run.

What’s next

I’d like to test whether the changes that helped here still help as the model grows.

A model above 1B parameters would give me a useful next comparison for the U-Net connections. Depending on the batch size and memory requirements, I may also need to explore tensor or pipeline parallelism.

I also want to compare tokenizers. The current model uses basic BPE; tiktoken’s byte-level BPE is one implementation to examine. I’d want to compare sequence lengths and downstream results before changing the training setup.

Mixture of experts is another possibility. Replacing some feedforward blocks with expert layers could add capacity without activating all of it for every token, though it would also add routing and training complexity.

I’d also like to extend low-precision training beyond the attention path used here. Work such as FP8-LM explores using FP8 more broadly during training. I haven’t tested that approach yet.

Making room for more experiments

The U-Net connections gave me an architectural idea that I could test without building a completely different model. Some variations helped in these experiments and others didn’t, so I want to keep comparing them against a plain Transformer as the setup changes.

Training speed mattered because it determined how many comparisons I could afford to run. Mixed precision, compilation, and checkpointing made the architectural experiments easier to carry out on the hardware I had.

For this project, it made sense to get a stable baseline and measure the runtime changes before spending more time on architectural variations. That gave me a better chance of telling whether a result came from the model itself, the training settings, or the implementation.

Code is on GitHub.