{"slug": "fun-with-low-rank-vocab-matrices-and-a-bonus-test-loss-reduction", "title": "Fun with low-rank vocab matrices (and a bonus test loss reduction?)", "summary": "Training GPT-2 small-style models with low-rank factorised embeddings on the input side alone reduced test loss, while applying the trick to both embeddings and the output head raised loss by 0.07, 0.10 and 0.14 across three tests, according to an experiment by the model's creator. The author was prompted by Hugging Face user AndrewThompson1233, who reported that at rank 128 on a 50k vocabulary the validation cross-entropy penalty is typically within +0.02 to +0.04 loss, or under 0.5 perplexity. The finding matters because embeddings and the output head account for roughly 39 million of the 163 million parameters in the author's GPT-2 small-style models, about half the model, or 30% under weight tying.", "body_md": "## Fun with low-rank vocab matrices (and a bonus test loss reduction?)\n\nI was [nerdsniped](https://xkcd.com/356/)! On the HF discussion page for one of the\nmodels I created for my [previous post](https://www.gilesthomas.com/2026/10/why-do-openai-gpt2-weights-beat-mine-5-data-quality),\n[`AndrewThompson1233`](https://huggingface.co/AndrewThompson1233)\nasked if I'd considered trying out [factorised embeddings](https://huggingface.co/gpjt/jax-with-mha-bias-fw-fwedu-5050-DEPRECATED/discussions/1)\n-- something he's using for his model, Maba, and which was previously used in some other models, like\n[ALBERT](https://arxiv.org/abs/1909.11942v6).\n\nIt's a really nifty idea; you use a similar trick to LoRA as a way of reducing the number of parameters used for your embeddings and your output head. And for small models, those can be a disproportionate number of the total.\n\nFor example, with my 163 million parameter GPT-2 small-style models, the embeddings\nare about [39 million of them](https://www.gilesthomas.com/2026/07/llm-parameter-counts); the output head is\nthe same size, so with those two taken together,\nthat's half the model just getting stuff in and out rather than actually doing the\nthinking. Even if you use weight tying, like the original GPT-2 models did, you'll\nwind up spending 30% of your \"budget\" on the single matrix shared between embeddings\nand the output head.\n\nSimilarly, while larger models have a smaller percentage spent that way, even medium-sized MoE models\ncan wind up spending a lot of the *active*\nparameters on embeddings; I [calculated](https://www.gilesthomas.com/2026/07/benchmarking-qwen-3-6-35b-moe-rtx-3090)\nthat for Qwen 3.6 35B MoE, with 3B active parameters, there were about 1B of them used\nacross both the input embeddings and the output head.\n\nSo anything that can reduce that -- so long as it doesn't come at a high cost in terms of the model's capabilities -- is worth considering. You save on your parameter count, which either means smaller models and potentially quicker training, or allows you to \"invest\" the savings in more thinking parameters -- that is, a wider network with a higher embedding dimensionality, or a deeper one with more layers.\n\nAndrew reported that the impact of using this trick was minimal:\n\nAt rank 128 on a 50k vocabulary, the penalty on validation cross-entropy is negligible (typically within +0.02 to +0.04 loss, or <0.5 perplexity delta).\n\nNow, that suggests that he's getting a loss of around 2.5 (that being the point where a loss increase of 0.04 implies a perplexity change of 0.5), which is much lower than I typically get with my GPT-2 architecture (3.5 is more like it for me), but still, perhaps I'd get good numbers too? I decided to train a few models to see what happened.\n\nThe results were interesting! I found that using this trick on both the embeddings and the output head caused a somewhat larger increase in loss than Andrew described -- about 0.07 in my first test, 0.10 in the second, and 0.14 in the third. Those are relatively large, though potentially recoverable from further training or larger models.\n\nBut more surprisingly, I found that using the trick on the input embeddings only\nseemed to *reduce* test loss. Before I'd started, I'd expected it to be less harmful\non the input side than it was on the output side, but seeing an improvement was certainly\nunexpected. Getting loss down by reducing your number of parameters is not what normally\nhappens!\n\nSo, it's definitely worth digging in a bit. Let's start with the theory: what are these factorised embeddings, and how do they work? Like LoRA, they rely on low-rank factors, so firstly I'll define those.\n\n### Low-rank factors\n\nOur models are made up of a large set of matrices -- embeddings, output heads, attention weights, FFN linear layers, and so on. Making them smaller obviously reduces the number of parameters for the model, though you would expect it to come at a cost. The idea behind low-rank factors is that you can replace a larger matrix with two smaller ones that will contain most of the important information.\n\nImagine that you've got a matrix of size . Now, from matrix multiplication,\nwe know that if you multiply an matrix by an one, you'll get\na result that is . If we do that, then we've factorised the matrix (in the\nsame way as we might factorise 12 into 3 and 4 because ), and the value is referred to (logically enough) as the *rank* of the factorisation.\nIf is sufficiently smaller than and , then\nthese two matrices -- the *low-rank factors* -- combined will contain fewer parameters than the full matrix.\n\nLet's make this specific, using the output head of a GPT-2 style model without weight tying. It takes in the embeddings that result from the Transformer layers (after normalisation), and converts them into logits across the vocabulary. For the GPT-2 small size, our embeddings are 768-dimensional, and the vocab size is 50,257 tokens. So when we use a linear layer to do that mapping, it has parameters.\n\nIf we were to replace it with two matrices -- say, one of and one of , then combined they would take up parameters. That is almost six times smaller!\n\nSo: the idea is that instead of creating our model with an matrix, we create it with a pair, and , and train those instead. If all goes well, that pair will be almost as capable of learning what we want as the full one would have been, so we'll get results that are close enough.\n\nThe maths and the implementation work out simply, too. For a linear layer with no bias,\nwe can write the matrix multiplication that takes our inputs and a weight matrix , and\nproduces an output , like this: [1](#fn-1)\n\nNow, if we're using a pair of low-rank factor matrices and instead of a full one, , we can say the \"virtual\" weights we want to use for the calculation are such that:\n\nSo that means that our calculation for this neural network is this:\n\nMatrix multiplication is associative, which means that you can rewrite that as:\n\n...which hopefully you can see is the same as feeding through a layer using as its weights, and then feeding the result through a second one using . In code, you've taken something like this:\n\n```\n out_head = nn.Linear(d_emb, vocab_size, bias=False)\n ...\n logits = out_head(embeddings)\n```\n\n...and replaced it with this:\n\n```\n out_head_A = nn.Linear(d_emb, r, bias=False)\n out_head_B = nn.Linear(r, vocab_size, bias=False)\n ...\n intermediate = out_head_A(embeddings)\n logits = out_head_B(intermediate)\n```\n\nIt's probably intuitively obvious, though, that this comes at a cost. If you have\nfewer parameters then you can store less information, so this part of your model\nis \"dumber\". If it's not clear, though, imagine that `r` in the code above was one.\nYou would be taking the (for GPT-2 small) 768-dimensional embedding, converting it\nto a single number, and then expanding that single number out to a 50,257-dimensional\nset of logits. Stuff is going to get lost -- the idea behind low-rank factors is that,\nso long as you choose an appropriate value for (which is often referred to as the\nwidth of the *low-rank bottleneck*), you won't lose the *important* stuff.\n\nThis works surprisingly well in many cases -- in particular, LoRA, which allows you to\nfine-tune models that are too hard to fully train on your hardware, or to fine-tune them faster,\nuses it to easily train something like low rank factor \"diffs\" to your weight matrices [2](#fn-2).\n\nBut for this post, we'll try the simplest version: what if we replace both the input embeddings and the output head -- the vocabulary matrices -- in their entirety with low-rank bottlenecks?\n\n### Introducing LoRE\n\nLow-rank factorisation of the vocabulary matrices is a bit of a mouthful, so based on the name \"LoRA\", I decided to call this trick LoRE, for low-rank embeddings. It's not strictly accurate (after all, we're doing it both to the embedding matrix at the start of the LLM and to the output head at the end), and I don't expect it to take off, but I rather like it and will use it in this post :-)\n\nThe argument behind it is that the embeddings are just a simple lookup table, so are\nexactly the kind of place you'd expect to be able to make savings by only considering\nthe important parts of our huge matrix; the same *might* apply to the output head at\nthe end, though personally I was less convinced by this part.\n\nIt's worth looking into that asymmetry a bit. My intuition was that while embeddings\nreally are a lookup table, the output head is doing something a bit more subtle. It's\nprojecting from the continuous embeddings that come out of our Transformer layers into\nlogits, which we interpret (via softmax) as a probability distribution over possible next\ntokens. That's a significantly less simple job than mapping from \"cat\" to\nthe appropriate embedding: as a (simplified) example, if you have an embedding that means\nsomething adjacent to \"cat\", \"dog\", \"gerbil\", \"household pet\" and so on, then projecting\nthat to the appropriate values for the next token is non-trivial. [3](#fn-3)\n\nStill, implementation-wise, it was really simple. I decided to extend the PyTorch code that\nI had from [Sebastian Raschka](https://sebastianraschka.com/)'s book\n\"[Build a Large Language Model (from Scratch)](https://www.manning.com/books/build-a-large-language-model-from-scratch)\",\nby adding an optional `lore` section to the model config JSON. This would allow us\nto switch on LoRE mode for the input embeddings and the output head independently,\nand would specify the rank -- that is, the in the section above, which says how wide\nthe low-rank bottleneck between matrices and would be.\n\nSo, in code, my first cut changed this:\n\n```\n self.tok_emb = nn.Embedding(cfg[\"vocab_size\"], cfg[\"emb_dim\"])\n```\n\n...to this:\n\n```\n if \"lore\" in cfg and cfg[\"lore\"].get(\"input_embeddings\", False):\n self.tok_emb = nn.Sequential(\n nn.Embedding(cfg[\"vocab_size\"], cfg[\"lore\"][\"rank\"]),\n nn.Linear(cfg[\"lore\"][\"rank\"], cfg[\"emb_dim\"], bias=False)\n )\n else:\n self.tok_emb = nn.Embedding(cfg[\"vocab_size\"], cfg[\"emb_dim\"])\n```\n\n...and this:\n\n```\n self.out_head = nn.Linear(\n cfg[\"emb_dim\"], cfg[\"vocab_size\"], bias=False\n )\n```\n\n...became this:\n\n```\n if \"lore\" in cfg and cfg[\"lore\"].get(\"output_head\", False):\n self.out_head = nn.Sequential(\n nn.Linear(\n cfg[\"emb_dim\"], cfg[\"lore\"][\"rank\"], bias=False\n ),\n nn.Linear(\n cfg[\"lore\"][\"rank\"], cfg[\"vocab_size\"], bias=False\n ),\n )\n else:\n self.out_head = nn.Linear(\n cfg[\"emb_dim\"], cfg[\"vocab_size\"], bias=False\n )\n```\n\nThat in itself was a perfectly reasonable implementation of the underlying concept. But it had a problem. I wanted to compare the loss with and without LoRE, and also find out what the effects of input embedding-only and output-head-only LoRE would be. But I've found in other experiments that the initial random weights that my models start with can have a significant effect on the loss that the resulting model gets. The code above would give quite different random weights with different configurations. The fix was simple, though:\n\n```\n if \"lore\" in cfg:\n lore_tok_emb = nn.Sequential(\n nn.Embedding(cfg[\"vocab_size\"], cfg[\"lore\"][\"rank\"]),\n nn.Linear(cfg[\"lore\"][\"rank\"], cfg[\"emb_dim\"], bias=False)\n )\n normal_tok_emb = nn.Embedding(cfg[\"vocab_size\"], cfg[\"emb_dim\"])\n if \"lore\" in cfg and cfg[\"lore\"].get(\"input_embeddings\", False):\n self.tok_emb = lore_tok_emb\n else:\n self.tok_emb = normal_tok_emb\n\n ...\n\n if \"lore\" in cfg:\n lore_out_head = nn.Sequential(\n nn.Linear(\n cfg[\"emb_dim\"], cfg[\"lore\"][\"rank\"], bias=False\n ),\n nn.Linear(\n cfg[\"lore\"][\"rank\"], cfg[\"vocab_size\"], bias=False\n ),\n )\n normal_out_head = nn.Linear(\n cfg[\"emb_dim\"], cfg[\"vocab_size\"], bias=False\n )\n if \"lore\" in cfg and cfg[\"lore\"].get(\"output_head\", False):\n self.out_head = lore_out_head\n else:\n self.out_head = normal_out_head\n```\n\nWith that code, so long as there is any LoRE config, we create the same weights\nregardless of whether it's used in the input embeddings or the output head. I would\nneed to train a baseline which had config for LoRE but just had both `input_embeddings` and\n`output_head` set to `False` there -- that would actually be a normal, non-LoRE model\nbut would have the same initial random weights in any parts that it shared with an\nequivalent model that did use LoRE.\n\nAdditionally, Andrew had [said](https://huggingface.co/gpjt/jax-with-mha-bias-fw-fwedu-5050-DEPRECATED/discussions/1#6aa983b11bcc35c715a81e94) that we need to be careful about initialisation:\n\nThe main failure mode to watch for during training is early gradient spikes: scaling the bottleneck projection initialization to 1 / sqrt(r) keeps variance uniform and ensures training dynamics match the full-rank baseline from step 1.\n\n...and [later](https://huggingface.co/gpjt/jax-with-mha-bias-fw-fwedu-5050-DEPRECATED/discussions/1#6aa9979704d7f855a3bbc3d5):\n\nKeep standard GPT-2 init on the embedding table, but scale the inner linear projection matrices by std = 1.0 / math.sqrt(128).\n\n(The 128 in there is because we were talking about a setup with , which is actually what I wound up using.)\n\nNow, I was using PyTorch's default initialisation rather than the classic GPT-2\nsetup, but I went ahead and added a `smart_initialize` flag to handle the scaling as\nAndrew described: if it was set, then right at the\nend of the model creation I applied the initialisation\nto the low rank factor matrices that were innermost in my model (as in, closest to\nthe Transformer layers):\n\n```\n if \"lore\" in cfg:\n if cfg[\"lore\"].get(\"smart_initialize\", False):\n nn.init.normal_(\n lore_tok_emb[1].weight,\n mean=0.0,\n std=1.0 / math.sqrt(cfg[\"lore\"][\"rank\"])\n )\n nn.init.normal_(\n lore_out_head[0].weight,\n mean=0.0,\n std=1.0 / math.sqrt(cfg[\"lore\"][\"rank\"])\n )\n```\n\nEven though that was not run for every configuration, because it happened right at the end of the model creation, it would not affect the other weights, so I figured that it would be safe.\n\nYou can see the full diff from the normal GPT-2 code [here](https://github.com/gpjt/ddp-base-model-from-scratch/compare/0d1aa53...d223121).\n\nThose changes looked correct to me, and I asked GPT-6 Astra (on \"extra high\" thinking)\nto confirm that it looked right to it in the context\nof the discussion, which it did. That done, I [double-checked with Andrew](https://huggingface.co/AndrewThompson1233/maba-v1-architecture/discussions/1),\nas I really wanted this to be a clean attempt at a repro of his results. He confirmed,\nso it was time to train some models!\n\n### The first training runs\n\nI decided that I wanted five training runs to check this out thoroughly:\n\n1. A baseline. As I said earlier, the LoRE stuff would have an effect on the randomness\nused for the initial weights, so I wanted to train one model with LoRE \"enabled\"\nbut not actually used. Config: [`model.json`](https://github.com/gpjt/ddp-base-model-from-scratch/blob/d2231213d6749921cf86f71e210d3f48da87f9da/runs/8xa100m40-lore-1-baseline-disabled/model.json) ,[`train.json`](https://github.com/gpjt/ddp-base-model-from-scratch/blob/d2231213d6749921cf86f71e210d3f48da87f9da/runs/8xa100m40-lore-1-baseline-disabled/train.json) .\n2. A model trained with LoRE enabled on both the input embeddings and the output\nhead, but without the smart initialisation of the low-rank matrices. This was\nsimply to confirm to myself that it was required. Config: [`model.json`](https://github.com/gpjt/ddp-base-model-from-scratch/blob/d2231213d6749921cf86f71e210d3f48da87f9da/runs/8xa100m40-lore-2-both-ordinary-init/model.json) ,[`train.json`](https://github.com/gpjt/ddp-base-model-from-scratch/blob/d2231213d6749921cf86f71e210d3f48da87f9da/runs/8xa100m40-lore-2-both-ordinary-init/train.json) .\n3. A model with LoRE on both ends, *with* the smart initialisation -- essentially,\nwhat Andrew had described. Config:[`model.json`](https://github.com/gpjt/ddp-base-model-from-scratch/blob/d2231213d6749921cf86f71e210d3f48da87f9da/runs/8xa100m40-lore-3-both-smart-init/model.json) ,[`train.json`](https://github.com/gpjt/ddp-base-model-from-scratch/blob/d2231213d6749921cf86f71e210d3f48da87f9da/runs/8xa100m40-lore-3-both-smart-init/train.json) .\n4. A model with LoRE on the input embeddings only, using smart initialisation if that had\nturned out to be necessary, and not using it if it had not. I deliberately made\nthe config invalid, using `TODO` instead of`true` or`false` for the`smart_initialize` flag, so that I'd be reminded to fill those in later.\nConfig:[`model.json`](https://github.com/gpjt/ddp-base-model-from-scratch/blob/d2231213d6749921cf86f71e210d3f48da87f9da/runs/8xa100m40-lore-4-input-only/model.json) ,[`train.json`](https://github.com/gpjt/ddp-base-model-from-scratch/blob/d2231213d6749921cf86f71e210d3f48da87f9da/runs/8xa100m40-lore-4-input-only/train.json) .\n5. A model with LoRE on the output head only, with smart initialisation the same as with (4).\nConfig: [`model.json`](https://github.com/gpjt/ddp-base-model-from-scratch/blob/d2231213d6749921cf86f71e210d3f48da87f9da/runs/8xa100m40-lore-5-output-only/model.json) ,[`train.json`](https://github.com/gpjt/ddp-base-model-from-scratch/blob/d2231213d6749921cf86f71e210d3f48da87f9da/runs/8xa100m40-lore-5-output-only/train.json) .\n\nTo save time, I decided to train them in the cloud: the first three, and then the last two,\nin two batches of parallel runs. I used 8x A100 machines with 40 GiB VRAM per GPU on [Lambda](https://lambda.ai/). I'll break\nwith my normal tradition of giving detailed runthroughs of each training run,\nas there are (spoiler) quite a lot of them in this post! However, for reproducibility,\nthe commands I ran for each one looked like this:\n\n```\ngit clone https://github.com/gpjt/ddp-base-model-from-scratch.git\ncd ddp-base-model-from-scratch/\ngit checkout lore-experiment\n./setup_lambda.sh\nuv sync\nuv run torchrun --nproc_per_node=8 ddp_train.py 8xa100m40-lore-1-baseline-disabled datasets\n```\n\n...with the run ID, `8xa100m40-lore-1-baseline-disabled`, replaced with the appropriate\none for the given [run settings](https://github.com/gpjt/ddp-base-model-from-scratch/tree/6ec817f0bbf525f65444574ab7b6a572a06d8461/runs).\n\nHere are the results from the first three runs: the baseline, and the ones with LoRE at both ends, with and without smart initialisation. The model name links go to the uploaded models on Hugging Face.\n\n| Model | Parameters | Training time (s) | Initial training loss | Final training loss | Test loss | Cost (US$) | \n|---|---|---|---|---|---|---|\n| [`8xa100m40-lore-1-baseline-disabled`](https://huggingface.co/gpjt/8xa100m40-lore-1-baseline-disabled) | 163,009,536 | 9,339 | 10.992 | 3.579 | 3.599611 | 43.96 | \n| [`8xa100m40-lore-2-both-ordinary-init`](https://huggingface.co/gpjt/8xa100m40-lore-2-both-ordinary-init) | 98,877,184 | 8,057 | 10.89 | 3.703 | 3.723328 | 38.55 | \n| [`8xa100m40-lore-3-both-smart-init`](https://huggingface.co/gpjt/8xa100m40-lore-3-both-smart-init) | 98,877,184 | est. 8,100 | 11.865 | unknown | 3.671585 | 41.44 | \n\nYou can see that I had some issues with the last run; there were shortages of available\ninstances, and I wound up starting one in Japan. Connectivity to it from my home in\nPortugal was really slow and unreliable; when I tried to download the model after the training run,\n`scp` was predicting 10 hours (which at the machine's cost of US$15.92/hour would have been insane). However,\nI was able to copy it to my [PythonAnywhere](https://www.pythonanywhere.com/) account\nquickly, perhaps because that is in a US Amazon datacenter and had better connectivity to Japan, and then I could copy it from\nthere to my machine, again quite quickly. Unfortunately while scrabbling around with that I didn't capture\nall of the output.\n\nStill, the important data at this stage was in the test loss column. As expected, using LoRE made the test loss worse. And smart initialisation did seem to help reduce that penalty.\n\nThe loss delta, at +0.07, was worse than the +0.02 to +0.04 that Andrew described in his message, but not wildly different -- certainly small enough that it seemed plausible that you could recover it, and perhaps more, by investing the 64M parameter difference in a deeper or wider model. Promising!\n\nThe cost and time savings were real, too: the\nrun in Japan had its cost bumped up by all of the woes I had copying the model down,\nso the `8xa100m40-lore-2-both-ordinary-init` one is the best indication: US$38.55 rather than\nUS$43.96, so a saving of more than 10%. And while going down from a bit more than 2h30m\nto 2h14m might not seem like much, with a larger model or a longer training run, an almost 14% reduction\ncould be valuable in and of itself. (And that's not even considering the alternative of\n\"reinvesting\" the saved parameters in larger models.)\n\nSo that was pretty good news! It was time to do the ablation runs to see whether\nhaving LoRE just on the token embeddings, or just on the output head, made things\nany different. I switched the `smart_initialize` setting for the remaining two\nmodels to `true`, and kicked them off.\n\n| Model | Parameters | Training time (s) | Initial training loss | Final training loss | Test loss | Cost (US$) | \n|---|---|---|---|---|---|---|\n| [`8xa100m40-lore-4-input-only`](https://huggingface.co/gpjt/8xa100m40-lore-4-input-only) | 130,943,360 | 9,225 | 10.978 | 3.505 | 3.526610 | 44.83 | \n| [`8xa100m40-lore-5-output-only`](https://huggingface.co/gpjt/8xa100m40-lore-5-output-only) | 130,943,360 | 8,202 | 11.817 | 3.717 | 3.737128 | 41.59 | \n\nNow, there's a bunch of interesting stuff in there, but let's focus on the test\nloss. My lab notes for this bit say \"well, that's a bit of a shocker!\"\nQuite amazingly -- at least to me -- we got *lower* test loss with the input-only LoRE,\nwith 3.526610, than we did with the no-LoRE baseline, which got 3.599611.\n\nThings get even more interesting when you compare the different versions (excluding the non-smart-init) against each other, in terms of their difference from the no-LoRE baseline (which I'll call the delta for conciseness below). Rounding to 3dps:\n\n| | Output LoRE off | Output LoRE on | \n|---|---|---|\n| Input LoRE off | - | 0.138 | \n| Input LoRE on | -0.073 | 0.072 | \n\nIf you add the two single-LoRE loss deltas together -- the 0.138 and the -0.073 -- you get 0.065, which is surprisingly close to the delta of 0.072 for the dual-LoRE option, just 0.007 away.\n\nCould we be looking at two almost independent effects, one for each option of where to apply the LoRE? And was that apparent benefit of the input-only LoRE a real thing? Or was it due to chance -- perhaps we had \"good luck\" when initialising the LoRE weights, and \"bad luck\" when initialising the non-LoRE ones. If there was an overlap between the best loss you could get with LoRE and the worst loss you could get without, then this result could be within the noise.\n\nI decided I was going to do a second batch of training runs to see if I could get evidence one way or the other. But before we get into that, it's worth taking a look at some of the other numbers in those two tables.\n\nFirstly, let's take a look at the training time. The training machines looked like they were running at close to 100% during all of the runs, so tentatively let's assume that time taken is proportional to compute. Unsurprisingly, the baseline run came in with the longest time, at 9,339s. And equally unsurprisingly, the full-LoRE run that we have timing for was much faster, at 8,057s. It certainly makes sense that replacing big matrices with pairs of (much) smaller low-rank factors will mean fewer calculations on both the forward and the backward pass, meaning that the training run is faster.\n\nThere was something that initially surprised me with the single-LoRE runs. Using LoRE on the output head sped things up quite a lot, but input embedding LoRE had much less effect.\n\nHowever, after thinking about it a bit, it became clear what was going on. An embedding\nmodule in a network is actually really quick to run, both on the forward and backward pass.\nWhile I like to think of embeddings as *conceptually* being a projection of a one-hot\nvector in vocab space into embedding space via a matrix multiplication, and I still think\nthat's an excellent model of interpreting what's going on, in practice it's just something\nmuch more like an array lookup, with almost zero cost. Likewise on the backward pass,\nwe only need to work out gradients for the chosen embedding, so that should be faster too.\n\nIndeed, with that all in mind, it might not have been surprising if input-only LoRE\nhad turned out to be a bit *slower* than the baseline, as it, at least, needed to do\nforward and backward passes through actual matrices, even though they were fairly small.\n\nDoing another of those delta tables, the time difference in seconds (using the \"ordinary init\" run for the both-on number):\n\n| | Output LoRE off | Output LoRE on | \n|---|---|---|\n| Input LoRE off | - | -1,137 | \n| Input LoRE on | -114 | -1,282 | \n\nYou can see that if you add the delta of the single-LoRE options, you get -1,251, which is pleasingly close to the delta of the dual-LoRE one. And that is much less surprising than the loss result -- with training compute, you really would expect improvements like this to be additive.\n\nFinally, there was something interesting in the initial training loss. Because my\ntraining code prints out the loss it got on the first global step, I decided to note it down.\nThere's something in this data that will become important later: in both of the training\nruns with LoRE on the output head *and* smart initialisation, it went up noticeably.\n\nNormally, the initial training loss on my GPT-2-style models is a bit less than 11. This fits in well with the architecture; you'd expect that a model that was predicting tokens purely randomly (well, strictly, uniformly) would have a perplexity roughly equal to the vocab size; the vocab size for the GPT-2 tokeniser is 50,257, and loss is the natural logarithm of perplexity, so that implies a loss of roughly 10.82.\n\nBut both of those smart-init models with LoRE on the output head had an initial loss about\none point higher, at 11.865 and 11.817, meaning that they were significantly *worse* than\nrandom. What might be going on here? We'll find out later :-)\n\nAnyway, for now, that was an interesting set of results. I wanted to dig in more; in particular, would that surprising improvement of test loss for the input-only LoRE model reproduce if I tried again with a different random seed?\n\n### The second set of models\n\nI had two working hypotheses at this point; either\n\n1. What we were seeing was due to luck of the draw with the random initialisation of the weights -- good luck with LoRE, bad luck without.\n2. There really was an effect here: input-only LoRE actually improved test loss in the resulting model.\n\nA solid and simple test to try to distinguish between these was just to train another set of models, with a different random seed. If I got the same kind of results, input LoRE beating non-LoRE, that would be evidence in favour of (2) rather than (1), whereas if things changed, it would work the other way.\n\nI decided not to bother with a non-smart-init model this time, so I set things up for\nfour new models. My training code [uses a random seed of 42](https://github.com/gpjt/ddp-base-model-from-scratch/blob/d2231213d6749921cf86f71e210d3f48da87f9da/ddp_train.py#L762)\nby default, but allows it to be overridden in the `train.json` config file, so I decided\nto use 123 for these, and created:\n\n- A new baseline with no LoRE. Config: [`model.json`](https://github.com/gpjt/ddp-base-model-from-scratch/blob/804cb3f60aba5157ce3ded2c3295de653bca6526/runs/8xa100m40-lore-6-new-seed-baseline-disabled/model.json) ,[`train.json`](https://github.com/gpjt/ddp-base-model-from-scratch/blob/804cb3f60aba5157ce3ded2c3295de653bca6526/runs/8xa100m40-lore-6-new-seed-baseline-disabled/train.json) .\n- An input-only LoRE. I wanted to run this one as the first non-baseline one, because\nif it got *worse* (higher) loss than the baseline, then it would look rather\nlike the effect I'd previously seen was due to random weight initialisation,\nso I might want to consider stopping there.\nConfig:[`model.json`](https://github.com/gpjt/ddp-base-model-from-scratch/blob/804cb3f60aba5157ce3ded2c3295de653bca6526/runs/8xa100m40-lore-7-new-seed-input-only/model.json) ,[`train.json`](https://github.com/gpjt/ddp-base-model-from-scratch/blob/804cb3f60aba5157ce3ded2c3295de653bca6526/runs/8xa100m40-lore-7-new-seed-input-only/train.json) .\n- If I decided to continue, the next one would be LoRE on both the input and output sides.\nConfig: [`model.json`](https://github.com/gpjt/ddp-base-model-from-scratch/blob/804cb3f60aba5157ce3ded2c3295de653bca6526/runs/8xa100m40-lore-8-new-seed-both-smart-init/model.json) ,[`train.json`](https://github.com/gpjt/ddp-base-model-from-scratch/blob/804cb3f60aba5157ce3ded2c3295de653bca6526/runs/8xa100m40-lore-8-new-seed-both-smart-init/train.json) .\n- Finally, I'd do an output head LoRE only one. Config: [`model.json`](https://github.com/gpjt/ddp-base-model-from-scratch/blob/804cb3f60aba5157ce3ded2c3295de653bca6526/runs/8xa100m40-lore-9-new-seed-output-only/model.json) ,[`train.json`](https://github.com/gpjt/ddp-base-model-from-scratch/blob/804cb3f60aba5157ce3ded2c3295de653bca6526/runs/8xa100m40-lore-9-new-seed-output-only/train.json) .\n\nOnce I had the config together, it was time to train the models -- but it was hard to\nget available instances on Lambda. I wound up having to use my [`lambda-manager`](https://www.gilesthomas.com/2026/04/automating-starting-lambda-instances)\nscript to start things up, but after a few hours I got an alert on my phone that I had\ninstances running, and could kick things off.\n\nThe second training run with input-only LoRE came in with a test loss that once again beat the baseline, so I did the full set of runs, and here are the results:\n\n| Model | Parameters | Training time (s) | Initial training loss | Final training loss | Test loss | Cost (US$) | \n|---|---|---|---|---|---|---|\n| [`8xa100m40-lore-6-new-seed-baseline-disabled`](https://huggingface.co/gpjt/8xa100m40-lore-6-new-seed-baseline-disabled) | 163,009,536 | 9,409 | 11.006 | 3.577 | 3.597584 | 45.03 | \n| [`8xa100m40-lore-7-new-seed-input-only`](https://huggingface.co/gpjt/8xa100m40-lore-7-new-seed-input-only) | 130,943,360 | 9,163 | 10.989 | 3.563 | 3.584511 | 43.03 | \n| [`8xa100m40-lore-8-new-seed-both-smart-init`](https://huggingface.co/gpjt/8xa100m40-lore-8-new-seed-both-smart-init) | 98,877,184 | 8,102 | 11.829 | 3.678 | 3.698644 | 37.91 | \n| [`8xa100m40-lore-9-new-seed-output-only`](https://huggingface.co/gpjt/8xa100m40-lore-9-new-seed-output-only) | 130,943,360 | 8,192 | 11.806 | 3.693 | 3.712537 | 38.32 | \n\nAnother matrix of deltas against the baseline:\n\n| | Output LoRE off | Output LoRE on | \n|---|---|---|\n| Input LoRE off | - | 0.115 | \n| Input LoRE on | -0.013 | 0.101 | \n\nSo, once again, input-only LoRE improved test loss (though at -0.013, it was a smaller effect than the -0.073 the previous set of models had).\n\nAnd once again, the sum of the input-only and output-only LoRE deltas was close to the combined-LoRE model's -- indeed it was even closer, 0.102 being just 0.001 away from 0.101.\n\nThe training time results and the higher initial training loss for models with LoRE on the output head also held up, which was good to see. So at this point it was time to wrap things up -- or was it?\n\n### The `smart_initialize` problem\n\nWhile I was doing the second batch of training runs, I started writing up this\nexperiment. As I [normally do](https://www.gilesthomas.com/2026/07/ai-use), I ran a draft past various AI\nmodels, and Claude Opus 5.5 on max thinking spotted something.\n\nAndrew had [written](https://huggingface.co/gpjt/jax-with-mha-bias-fw-fwedu-5050-DEPRECATED/discussions/1#6aa9979704d7f855a3bbc3d5):\n\nKeep standard GPT-2 init on the embedding table, but scale the inner linear projection matrices by std = 1.0 / math.sqrt(128).\n\nI'd interpreted that \"inner\" as meaning that the \"innermost\" matrices in our low rank factors -- the ones that were closest to the Transformer layers -- should be initialised as he said, with values drawn from a normal distribution with a standard deviation of , and a mean of zero.\n\nSo that implied the second matrix in the embeddings, and the first one in the output head:\n\n```\n if \"lore\" in cfg:\n if cfg[\"lore\"].get(\"smart_initialize\", False):\n nn.init.normal_(\n lore_tok_emb[1].weight,\n mean=0.0,\n std=1.0 / math.sqrt(cfg[\"lore\"][\"rank\"])\n )\n nn.init.normal_(\n lore_out_head[0].weight,\n mean=0.0,\n std=1.0 / math.sqrt(cfg[\"lore\"][\"rank\"])\n )\n```\n\nClaude felt that it should be the second low-rank factor in both cases -- that is, we should change\n\n```\n lore_out_head[0].weight,\n```\n\n...to\n\n```\n lore_out_head[1].weight,\n```\n\nIt made a very good case for doing things that way, and digging into that is useful because it changes this \"smart\" initialisation from being a blindly-implemented \"secret sauce\" into something actually meaningful.\n\nLet's imagine that we're not using LoRE, and consider the output head only. It is\na PyTorch [`nn.Linear`](https://docs.pytorch.org/docs/2.14/generated/torch.nn.Linear.html),\nwhich means that:\n\n...the values are initialized from , where\n\nSo, we have a matrix, which we can call , that is full of numbers drawn from that probability distribution -- a uniform (flat) distribution within the range from to .\n\nFor LoRE, we're replacing with two separate matrices, which we'll call and :\n\n```\n self.out_head = nn.Sequential(\n nn.Linear(\n cfg[\"emb_dim\"], cfg[\"lore\"][\"rank\"], bias=False\n ),\n nn.Linear(\n cfg[\"lore\"][\"rank\"], cfg[\"vocab_size\"], bias=False\n ),\n )\n```\n\nLet's say that is shaped , which means that is , and is . We want the combined effect of those two matrices to start off with random numbers that look like -- in terms of their randomness -- the full matrix .\n\nBut what do the numbers actually look like in our \"virtual\" matrix with the code above? On the face of it, it sounds unlikely to be anything like the probability distribution for . After all, was initialised based on a formula that used . You can see that will have been initialised using the same distribution, as it has the same number of input features, but the number of input features for is , so its random numbers will have been based on that instead. Their combination into will presumably blend something from each of those two different probability distributions.\n\nWe can actually work out what that will be. Maths incoming: [click here to skip](#correcting-the-smart-initialisation).\n\nPer [Wikipedia](https://en.wikipedia.org/wiki/Distribution_of_the_product_of_two_random_variables),\nthe variance of the product of two independent random scalar variables (so, not matrices, but\nwe'll come back to that) and is this:\n\n...where is 's variance, is its mean, and likewise for .\n\nNow, let's say that the mean is zero for both and (like it is for PyTorch's distribution for the initial weights, ). That simplifies the above:\n\nHow does that apply to our matrices?\n\nWell, let's write out how we calculate by multiplying and , assuming that , , and :\n\nBy normal matrix maths, the value at position in the result will be the dot product of the zeroth row in the first matrix (taken as a vector) and the zeroth column in the second:\n\nWe can rewrite that for an arbitrary low rank bottleneck like this:\n\n...and even more generally, we can define all of the items in like this:\n\nSo, every element is the sum of products, where each product multiplies a number from whatever distribution was used to initialise 's elements by a number from whatever was used for 's.\n\nUsing the notation from above, 's elements have variance and 's have variance . So that means that the variance of each component of that sum is . There are of them; if you add together independent variables, their variances add, so that means that the variance of the elements of our resulting is:\n\nNow, let's remember that our matrix actually already was initialised in the way\nwe would have liked our non-LoRE version to have been done. The distribution of the random initial weights for an `nn.Linear` is\ndependent entirely on the number of input features -- and of course has the same number\nof input features as .\n\nSo how can we make be ? We need to make , and then the equation above drops out like this:\n\nNow, let's look at the code (with Claude's correction in there):\n\n```\n nn.init.normal_(\n lore_out_head[1].weight,\n mean=0.0,\n std=1.0 / math.sqrt(cfg[\"lore\"][\"rank\"])\n )\n```\n\nWe're setting the weights on our matrix to be drawn from a distribution with a standard deviation of . SD is, of course,\nthe square root of the variance (which is why we've been using things\nlike with that squaring in there to represent variance -- on its\nown is used for standard deviation). So the variance of the matrix after we've called\nthat `normal_` on it will be as desired, and will have random initial weights\nwith the same variance as would have done if we'd created it directly.\n\nThat's rather satisfying :-)\n\nSo that's the output head. How about the input embeddings? They're\n[`nn.Embedding`](https://docs.pytorch.org/docs/2.14/generated/torch.nn.Embedding.html)s,\nwhich are initialised somewhat differently:\n\ninitialized from\n\nThat's a normal distribution, with a mean of 0 and a variance (and thus also a standard deviation) of 1. The maths above applied to any distribution with a mean of zero -- so the smart initialisation should work for those too. Our matrix would have the correct variance, 1, and a mean of zero, and therefore so would , given that (for embeddings) already had the smart initialisation.\n\nSo does that mean that our \"virtual\" matrix created by multiplying and is identical in terms of randomness to the matrix that it's replacing? Well, not quite.\n\nRemember that the output head had numbers drawn from a uniform distribution, and the embeddings from a normal distribution. We're multiplying the matrices that are the \"input\" low rank factor matrices by s that are initialised from a normal distribution. The variance of the matrices is maintained by the smart initialisation, as we've shown, and the mean (being zero) is going to come through cleanly too. But the shape will change.\n\nI did a bit of digging around and it started to feel like a bit of a rabbit hole, but the one thing I was able to be certain of is that a variable drawn from a uniform distribution multiplied by another from a normal distribution gives a result that is neither uniform nor normal, and that normal times normal gives something called a normal product distribution (or also a product-normal distribution), which is also neither uniform nor normal.\n\nRemember that each element of our matrix is defined by this:\n\nSo each one is a sum of numbers, each of which is drawn from a product normal distribution\n(for embeddings) or from whatever the uniform normal distribution is called. What shape\nwould that be? Again: rabbit hole, but the central limit theorem states ([per Wolfram MathWorld](https://mathworld.wolfram.com/CentralLimitTheorem.html)):\n\nthe normalized sum of independent random variates with finite variances approaches a normal distribution\n\nSo, for large , we can say that the result is going to be close to a normal distribution. [4](#fn-4)\n\nWhat the smart initialisation does is make sure that our \"virtual\" weights, , have the same mean and variance as the ones we're trying to approximate, . But it won't keep the same shape for the distribution that they're drawn from -- certainly not for the output head, which goes from uniform to normal-ish. The embeddings are less impacted (though not completely untouched) because they were normal and are still almost so.\n\nStill, it's a neat trick! Was it what Andrew had actually meant?\n\n### Correcting the smart initialisation\n\nHaving convinced myself that the smart initialisation code I'd been using was wrong, I decided to check with Andrew. My first thought was that perhaps he was using weight tying -- remember, with weight tying, the embedding matrix is just re-used in transposed form for the output head. That would mean that the embeddings would be , and you'd use the smart initialisation on , but the output head -- being that transpose -- would use for the same and , so the smart initialisation would be the other way around.\n\nBut as it turned out, it was [just a miscommunication](https://huggingface.co/AndrewThompson1233/maba-v1-architecture/discussions/1#6abfe78119024613efda8f02).\nHe said that he keeps his weights untied, and had been using the word\n\"inner\" in \"inner linear projection matrices\" to mean something different to what I'd\ninterpreted it as. He confirmed that the real smart init code should always be\napplied to the second matrix, as the maths (and Claude) suggested.\n\nSo that was a simple change. You can see [the diff here](https://github.com/gpjt/ddp-base-model-from-scratch/commit/6ec817f0bbf525f65444574ab7b6a572a06d8461),\nbut in short:\n\n- I changed the `smart_initialize` config parameter so that instead of just being`False` or`True` , it could be`False` ,`\"original\"` for the original incorrect\ninterpretation, or`\"corrected\"` for the fixed one.\n- I updated all of the existing training configurations that had `smart_initialize` set\nto`True` to be`\"original\"` .\n- I added on yet more config for some more training runs to test this, setting\n`smart_initialize` to the new`\"corrected\"` option.\n\nNow, the good news (both for my own sanity and my wallet) was that testing this fix\nwas not going to require another four training runs. My training code is essentially deterministic\n-- no dropout or any other randomness is used after the model is created [5](#fn-5).\nBecause the smart initialisation code was at the end of the model creation, and\nsmart initialisation of the output head was right at the end, that meant that changing\nit would not affect any other code that used randomness.\n\nSo if I did two training runs -- with LoRE switched on for the output head only, and with it switched on for both embeddings and output -- then those would be \"compatible\" with earlier training runs that had no LoRE at all, and that only had it on the embeddings. You can see the configs here (with no explicit random seed, so they'd use my default of 42):\n\n- Both: [`model.json`](https://github.com/gpjt/ddp-base-model-from-scratch/blob/6ec817f0bbf525f65444574ab7b6a572a06d8461/runs/8xa100m40-lore-10-both-corrected-smart-init/model.json) ,[`train.json`](https://github.com/gpjt/ddp-base-model-from-scratch/blob/6ec817f0bbf525f65444574ab7b6a572a06d8461/runs/8xa100m40-lore-10-both-corrected-smart-init/train.json) .\n- Output-only: [`model.json`](https://github.com/gpjt/ddp-base-model-from-scratch/blob/6ec817f0bbf525f65444574ab7b6a572a06d8461/runs/8xa100m40-lore-11-output-only-corrected-smart-init/model.json) ,[`train.json`](https://github.com/gpjt/ddp-base-model-from-scratch/blob/6ec817f0bbf525f65444574ab7b6a572a06d8461/runs/8xa100m40-lore-11-output-only-corrected-smart-init/train.json) .\n\nHere are the results; I've put them into a table with their counterparts from the original seed-42 training run:\n\n| Model | Parameters | Training time (s) | Initial training loss | Final training loss | Test loss | Cost (US$) | \n|---|---|---|---|---|---|---|\n| [`8xa100m40-lore-1-baseline-disabled`](https://huggingface.co/gpjt/8xa100m40-lore-1-baseline-disabled) | 163,009,536 | 9,339 | 10.992 | 3.579 | 3.599611 | 43.96 | \n| [`8xa100m40-lore-10-both-corrected-smart-init`](https://huggingface.co/gpjt/8xa100m40-lore-10-both-corrected-smart-init) | 98,877,184 | 8,069 | 10.987 | 3.722 | 3.741209 | 37.65 | \n| [`8xa100m40-lore-4-input-only`](https://huggingface.co/gpjt/8xa100m40-lore-4-input-only) | 130,943,360 | 9,225 | 10.978 | 3.505 | 3.526610 | 44.83 | \n| [`8xa100m40-lore-11-output-only-corrected-smart-init`](https://huggingface.co/gpjt/8xa100m40-lore-11-output-only-corrected-smart-init) | 130,943,360 | 8,194 | 10.986 | 3.766 | 3.785209 | 38.19 | \n\nBefore we look at the test loss, check out that \"Initial training loss\" column. The\nweirdness that I'd noticed earlier, where that value was much worse than uniform, is gone!\nThat actually makes some kind of sense. With the smart initialisation on the wrong\npart of the low-rank output head, our variance was completely wrong, and so the model\nwas initially producing particularly bad results. [6](#fn-6)\n\nSo that was promising! But the test loss results were even more interesting. In our original seed-42 results with the incorrect init, the output-head-LoRE-only model got a test loss of 3.737128, but here we got 3.785209. It was actually worse with the corrected smart initialisation! And that carried through to the result for models with LoRE on both input and output; the original model with the incorrect init got 3.671585, and here we got 3.741209 -- an even bigger worsening.\n\nCombining those two together, an image I like in terms of the loss landscape is that previously we were starting on a mountain that was near a deep valley -- we started with high loss, but there was a low-loss place nearby. But with the smart initialisation fixed, we were now starting on a hill with a less-deep valley nearby.\n\nSo what happens if we do one of our test loss delta tables showing how each model performed against the baseline?\n\n| | Output LoRE off | Output LoRE on | \n|---|---|---|\n| Input LoRE off | - | 0.186 | \n| Input LoRE on | -0.073 | 0.142 | \n\nIf we add up the effect of the input-only test loss delta of -0.073 and the output-only delta of 0.186, we get 0.113 -- quite different to the delta of 0.142 that we actually got with them both. That nice \"additive\" property that we had, where the delta from input LoRE only plus the delta from output LoRE only summed to almost exactly the same as the delta for LoRE on both ends, appears to have gone away :-(\n\nStill, an interesting set of results! Let's bring everything together.\n\n### The results\n\nHere are all of the training runs together in one table. I've removed the columns for the cost, the end training loss, and the training time to keep things manageable, and I've sorted it by the test loss.\n\n| Model | Parameters | Input LoRE | Output LoRE | Seed | Smart init | Initial training loss | Test loss | \n|---|---|---|---|---|---|---|---|\n| [`8xa100m40-lore-4-input-only`](https://huggingface.co/gpjt/8xa100m40-lore-4-input-only) | 130,943,360 | yes | no | 42 | original | 10.978 | 3.526610 | \n| [`8xa100m40-lore-7-new-seed-input-only`](https://huggingface.co/gpjt/8xa100m40-lore-7-new-seed-input-only) | 130,943,360 | yes | no | 123 | original | 10.989 | 3.584511 | \n| [`8xa100m40-lore-6-new-seed-baseline-disabled`](https://huggingface.co/gpjt/8xa100m40-lore-6-new-seed-baseline-disabled) | 163,009,536 | no | no | 123 | none | 11.006 | 3.597584 | \n| [`8xa100m40-lore-1-baseline-disabled`](https://huggingface.co/gpjt/8xa100m40-lore-1-baseline-disabled) | 163,009,536 | no | no | 42 | none | 10.992 | 3.599611 | \n| [`8xa100m40-lore-3-both-smart-init`](https://huggingface.co/gpjt/8xa100m40-lore-3-both-smart-init) | 98,877,184 | yes | yes | 42 | original | 11.865 | 3.671585 | \n| [`8xa100m40-lore-8-new-seed-both-smart-init`](https://huggingface.co/gpjt/8xa100m40-lore-8-new-seed-both-smart-init) | 98,877,184 | yes | yes | 123 | original | 11.829 | 3.698644 | \n| [`8xa100m40-lore-9-new-seed-output-only`](https://huggingface.co/gpjt/8xa100m40-lore-9-new-seed-output-only) | 130,943,360 | no | yes | 123 | original | 11.806 | 3.712537 | \n| [`8xa100m40-lore-2-both-ordinary-init`](https://huggingface.co/gpjt/8xa100m40-lore-2-both-ordinary-init) | 98,877,184 | yes | yes | 42 | none | 10.89 | 3.723328 | \n| [`8xa100m40-lore-5-output-only`](https://huggingface.co/gpjt/8xa100m40-lore-5-output-only) | 130,943,360 | no | yes | 42 | original | 11.817 | 3.737128 | \n| [`8xa100m40-lore-10-both-corrected-smart-init`](https://huggingface.co/gpjt/8xa100m40-lore-10-both-corrected-smart-init) | 98,877,184 | yes | yes | 42 | corrected | 10.987 | 3.741209 | \n| [`8xa100m40-lore-11-output-only-corrected-smart-init`](https://huggingface.co/gpjt/8xa100m40-lore-11-output-only-corrected-smart-init) | 130,943,360 | no | yes | 42 | corrected | 10.986 | 3.785209 | \n\nA few things stand out:\n\n- The input-only LoRE models lead the pack -- though it's a close-run thing,\nand the gap between `8xa100m40-lore-7-new-seed-input-only` and`8xa100m40-lore-6-new-seed-baseline-disabled` could well be[within the noise](https://www.gilesthomas.com/2026/04/llm-from-scratch-32i-interventions-what-is-in-the-noise) at 0.013. In that post, I trained three models with different weight initialisation seeds,\nand got models that differed by up to ~0.017.\n- The two models with the corrected smart initialisation are the worst -- even worse\nthan the model with no extra initialisation, `8xa100m40-lore-2-both-ordinary-init` .\n\nSo what does that all tell us?\n\n### Wrapping up\n\nThe numbers above suggest three things -- with important caveats below:\n\n1. LoRE can reduce the number of parameters in these models dramatically even if\nonly applied to the input embeddings, and that comes with potentially an\n*improvement* in test loss (0.073 in one test, 0.013 in a second).\n2. When applied to the output head, it has the same size of reduction in the parameter\ncount, and causes a *worsening* in test loss (over three tests,\n0.115, 0.138 and 0.186).\n3. The results with LoRE on both ends seemed to come in somewhere around (in two of the three cases, very close to) the sum of the input-only and output-only deltas, though it doesn't seem to be a strictly additive relationship.\n\nEven more caveated, we might say:\n\n1. The mathematically cleaner implementation of what I've been calling \"smart initialisation\" actually worked worse than the mistaken version I used originally, though this was only tested with one seed, so the evidence is much weaker.\n\nNeither the larger of the two apparent improvements from input-only LoRE nor the worsening from output-only\nwere small -- when I was trying various\n[interventions](https://www.gilesthomas.com/2026/04/llm-from-scratch-32j-interventions-trying-to-train-a-better-model-in-the-cloud#the-interventions) into\nmy GPT-2 models (like removing dropout, scheduling the learning rate, gradient clipping and\nso on), many of them had a similar or smaller level of effect. Similarly, training a comparable non-LoRE model on double the number\nof Chinchilla-optimal tokens (ie. on 40 per parameter rather than 20) improved loss\nby 0.09, while investing the same compute resources to do a Chinchilla-optimal training\nrun on a larger model improved it by roughly 0.13 over the baseline (see [this post](https://www.gilesthomas.com/2026/08/chinchilla-check)\nfor the raw numbers).\n\nSo, is input-only LoRE One Weird Trick To Improve Your Test Loss? Of course not. Are output-only LoRE and both-ends LoRE Bad Things? Not necessarily.\n\nThe first thing to remember is that this is a really limited set of experiments.\nAlthough I wound up spending about US$450 on training runs (I really should be investing\nin building my machine [`poppy`](https://www.gilesthomas.com/2026/07/poppy-the-training-box-1-the-beginnings) up as a better\ntraining box to save on Lambda costs), I was only able to test:\n\n- Against my own GPT-2 style model (important differences from the original: no weight tying or bias on the QKV weights, no dropout, and PyTorch default weight initialisation).\n- With two seeds.\n- With two options for smart initialisation, and one of them only with one seed.\n- Over a number of tokens chosen to be Chinchilla-optimal for the baseline non-LoRE model. [7](#fn-7)\n\nEven in order to say anything definitive about whether LoRE works for this specific architecture, it would be necessary to do tests over a much larger number of random seeds and different lengths (in terms of number of tokens) of training run:\n\n- Maybe the effect was just limited to the seeds I tested? Two is better than one, but three would be better still, four even better, and so on.\n- Maybe the effect disappears in longer training runs? Perhaps LoRE causes some kind of loss in capacity that stops models from improving past some point. Intuitively that could make sense -- if you imagine training as being in some sense compressing information from the training data into the parameters, having fewer parameters is clearly going to limit your capacity in some sense.\n\nTo get solid results on whether smart initialisation is a good thing, you'd need to try a number of seeds to see whether the poor performance of the \"correct\" smart initialisation against the \"incorrect\" one worked out as a general result. And if it did, why might that be?\n\nBeyond that: the apparent improvement in the loss from input-only LoRE is something\nthat came up in these experiments. Maybe it would turn out to be a real thing, maybe not,\nbut the real benefit that Andrew mentioned in his [original message](https://huggingface.co/gpjt/jax-with-mha-bias-fw-fwedu-5050-DEPRECATED/discussions/1#6aa94b7ff679b1d676c0a74a)\nwas about something quite different: LoRE saves a *lot* of parameters and only (as he\nsaid later) costs a small amount of loss. If you reinvest those saved parameters in\na deeper network with more layers, you can hopefully gain back the loss that you,\num, lost, and then some -- and wind up with a more capable model.\n\nSo important further tests (which I think I will do, but later on) would be:\n\n- Isoparameter tests:\n - Input-only LoRE saves 32M parameters, so what happens if you \"spend\" that on extra layers? Or alternatively make the model wider (in terms of its embedding dimension)? Or some combination of both?\n - Output-only LoRE likewise.\n - And, of course, both input and output.\n- Isocompute tests. Notably, input-only LoRE doesn't save you much compute, as I mentioned earlier -- embedding layers are cheap in terms of compute, if not in terms of parameters. So here perhaps it would be more interesting to see what happens if you did output-only LoRE, but trained the model on more data so that the total training compute budget matched. Would you gain back the lost loss?\n\nIt's worth noting that in my Chinchilla experiments, I scaled up the model by . That's about a 41% increase, which is quite close to the 39% we save by using LoRE on both ends. So if the loss improvements carried over, perhaps it might work? On the other hand, those models were also trained on 41% more tokens, so...\n\nSo there's a huge number of further experiments that would be needed to build out these results -- and that's even before we move beyond this specific architecture. Andrew reports solid results with his setup, and his model Maba does appear to be quite different. But would it extend to a Qwen-style one? Or Mistral?\n\nNext, there's what happens with scaling. LoRE is valuable in small models because embeddings and the output head are such a large proportion of the total parameter count. My guess is that it would become increasingly unhelpful as models got larger.\n\nTo take an extreme case, looking at [the Kimi K3 paper](https://arxiv.org/abs/2607.24653), it has a vocab size\nof 160k, and its hidden dimension (which I'll assume is what the embeddings are using)\nis 7,168. That means that the embedding layer uses something like 1,146,880,000 parameters, as\ndoes the output head, for a total of ~2.3B parameters (assuming they're not using weight tying,\nwhich seems likely for a monster like this).\nIt has 104.2B active parameters\nper token, so if it were to use LoRE you'd save something of the order of 2% of the active parameter\ncount. But for purposes of \"reinvesting\" in further layers, the real comparison point is the\ntotal number of parameters, not the active ones, and with about 2.8T total, the\nnumber used for embeddings is less than 0.1%. It doesn't sound\nlike it would help much.\n\nUnless, of course, the benefit of the input-only LoRE loss improvement actually turned out to be a real thing, and didn't weaken as the model was scaled...\n\nBut anyway, that's speculation. I think a good set of experiments for that would be to try the same training runs as before, but for larger models. Maybe just trying GPT-2 medium and large sizes would be a good start?\n\nLots of work. If I had access to Google's resources then I think I'd be kicking off a bunch of training runs. OTOH if I had access to Google's resources I probably wouldn't be blogging about this...\n\nI think that there's one thing I'm going to take away from this. LoRE is an interesting intervention -- and I think is well worth trying out on any new model where vocab matrices are a significant fraction of the parameters, just to see if it helps. I'll be doing that in future.\n\nAnd if anyone else out there wants to try LoRE-style training runs, I'd love to hear from you. Any results, positive or negative, would be really interesting.\n\nThanks for reading! And many, many thanks to Andrew for suggesting this as an idea to play around with -- it's been fun :-)\n\n1. \nThis is often written to reflect the fact that weights are generally kept in format, at least in PyTorch, but that's really [an implementation detail](https://www.gilesthomas.com/2025/02/basic-neural-network-matrix-maths-part-2) .[↩](#fnref-1)\n2. \nI'll write more about LoRA in the future -- I've been playing with it a bit -- but essentially, what you do is: \n - Freeze all of the matrices that make up your parameters (eg. in PyTorch, set\n`requires_grad` to`False` ).\n - For each matrix that you want to train (which might not be all of them -- for fine-tuning, people often only train a subset), you add on a *low-rank adapter* , which\nis a pair of matrices like our and above, and which sits \"beside\" the\nexisting one, like this:Gradient tracking is *enabled* on the matrices in the low-rank adapter, and\nits results are added to the results of the original weights.\n The idea is that the adapter pair learns to create a \"diff\" to the results of the original weights, which allows a certain amount of training of that part of your model. And it works surprisingly well. Again: I'll write about this in more depth in the future. [↩](#fnref-2)\n3. Freeze all of the matrices that make up your parameters (eg. in PyTorch, set\n4. \nWhen I was discussing this with an LLM -- sadly I've forgotten which -- the *softmax bottleneck* came up here, as it did in later correspondence with Andrew.\nThat's a reference to[this paper](https://arxiv.org/abs/1711.03953) . I'll need to\nwork through it carefully at some later point, but the relevant point appears to be\nthat the rank (which we can loosely read as how many independent ways the predictions can vary from one context to another) of the values coming out of softmax\nis limited by the rank of what is going into it. If we force our (in my GPT-2\nsmall case) 768-dimensional embeddings through a 128-dimensional bottleneck, we're\nreducing their inherent rank, and it's that reduced rank that will drive softmax's,\nso the results will be worse than what you'd get with the original 768. Seems logical enough.Linkrot-proof reference: Yang, Z., Dai, Z., Salakhutdinov, R., & Cohen, W. W. (2018). Breaking the softmax bottleneck: A high-rank RNN language model. In *International Conference on Learning Representations (ICLR)* .[↩](#fnref-3)\n5. \nThanks to Claude Opus 5.5 for pointing me at the central limit theorem. I definitely need to brush up on my statistics, and I'm sure I knew about it at some point back in the 90s... [↩](#fnref-4)\n6. \nModulo CUDA jitter of various flavours -- though even then, [a while back](https://www.gilesthomas.com/2026/04/llm-from-scratch-32i-interventions-what-is-in-the-noise) I did\ntwo training runs on 8x A100 machines with the same config, months apart, and\nthe Safetensors files were identical bit-for-bit. So I think the effect of that\nin my setup is minuscule.[↩](#fnref-5)\n7. \nThe temptation to dig into the details of why, and exactly what the mistaken smart initialisation did to the variance of the \"virtual\" weights , is strong. But I will resist. [↩](#fnref-6)\n8. \nWhen I ran a draft of this post past my [editorial board of LLMs](https://www.gilesthomas.com/2026/07/ai-use) ,\nGrok made the point that perhaps the apparent benefit of having LoRE\nonly on the input embeddings might have been due to that model seeing significantly\nmore tokens per parameter. It's an interesting thought! If we do the maths,\nthe input-only and output-only LoRE models were trained on ~25 tokens per parameter\nand the both-ends ones were trained on about 33, vs 20 for the non-LoRE ones.However, in my Chinchilla experiments, I only got 0.09 improvement in loss by training on 40 tokens per parameter rather than 20. Now, I'm sure that gain mostly happened early in the run -- overtraining is subject to diminishing returns -- but the result of 0.073 for input-only LoRE would imply that almost all of it happened in the first extra 5 tokens per parameter, which seems dubious. Those are different training runs with different seeds, and were using my JAX code, which has a better baseline loss than my PyTorch code, so we can't compare them directly. But it feels unlikely that the effect would be drastically different. Also, if we use *reductio ad absurdum* on Grok's logic, it does sound a bit weak --\nif we get rid of parameters but train on the same number of\ntokens, the higher tokens per parameter will not in general counteract the damage from the smaller model.\nImagine, for example, removing 32M parameters worth of Transformer layers and training\non the Chinchilla-optimal tokens for the size of the model before you did that. I'm not\ngoing to run the experiment, but I'd be... somewhat more than surprised if it improved the model.[↩](#fnref-7)\n\n## Citing this post\n\nThis is a blog, and if you want to link to this post then please do :-) However, if you're writing something more academic and need to do a proper citation, then here's a BibTeX block to make things easier.\n\n```\n@misc{thomas2026oct-low-rank-vocab-matrices,\n author = {Thomas, Giles},\n title = {{Fun with low-rank vocab matrices (and a bonus test loss reduction?)}},\n year = {2026},\n month = oct,\n howpublished = {Blog post},\n url = {https://www.gilesthomas.com/2026/10/low-rank-vocab-matrices},\n}\n```\n\n", "url": "https://wpnews.pro/news/fun-with-low-rank-vocab-matrices-and-a-bonus-test-loss-reduction", "canonical_source": "https://www.gilesthomas.com/2026/10/low-rank-vocab-matrices", "published_at": "2026-10-09 14:31:03+00:00", "updated_at": "2026-10-09 14:54:34.612858+00:00", "lang": "en", "topics": ["large-language-models", "machine-learning", "ai-research", "natural-language-processing"], "entities": ["GPT-2", "AndrewThompson1233", "Hugging Face", "ALBERT", "Qwen 3.6 35B MoE", "LoRA", "Maba"], "also_reported_by": [], "alternates": {"html": "https://wpnews.pro/news/fun-with-low-rank-vocab-matrices-and-a-bonus-test-loss-reduction", "markdown": "https://wpnews.pro/news/fun-with-low-rank-vocab-matrices-and-a-bonus-test-loss-reduction.md", "text": "https://wpnews.pro/news/fun-with-low-rank-vocab-matrices-and-a-bonus-test-loss-reduction.txt", "jsonld": "https://wpnews.pro/news/fun-with-low-rank-vocab-matrices-and-a-bonus-test-loss-reduction.jsonld"}}