New Record: Partition Residual Dims and Symmetric Value Embeddings (-0.7s)#194
New Record: Partition Residual Dims and Symmetric Value Embeddings (-0.7s)#194photomz wants to merge 4 commits into
Conversation
photomz
changed the title
|
Cool, hopefully this holds up when reran on the latest PR. Can you please drop the pushover file, and make sure train.py includes the updated step count? Been hoping for a while that someone could get the value embedding to work, maybe the addition of the value embed gates is what finally made it worthwhile, as it was not worthwhile in the past for some odd reason. When I added the ve_gates I did one test run with them on the last dimensions and got worse results so I kept them on the first 12. So that one surprises me a bit, glad you were able to get it to work. At the time the gates were in Muon and had much more volatile behavior... perhaps now that they are in Adam and not monopolizing the residual stream dimensions completely, they function better spread out, not sure. Splitting the dims with the skip gate is an interesting idea. |
|
I reran on #191 and confirmed the 15 step difference holds up. Retiming
Treatment # Treatment at 1745+40 (5x), timed at 1785 steps
import scipy.stats
import torch
losses = [
3.277315139770508,
3.276496648788452,
3.2764573097229004,
3.2788822650909424,
3.27580189704895,
]
times = [
109531.29593900303,
109799.04070798875,
109788.34327799996,
109647.2150749978,
109639.68391300295,
]
print("p=%.4f" % scipy.stats.ttest_1samp(losses, 3.28, alternative="less").pvalue)
# p=0.0024
print("losses:", torch.std_mean(torch.tensor(losses)))
# losses: (tensor(0.0012), tensor(3.2770))
print("time:", torch.std_mean(torch.tensor(times)))
# time: (tensor(110.0), tensor(109681.1))Control # Control at 1760+40 (5x), timed at 1800 steps
import scipy.stats
import torch
losses = [
3.277698040008545,
3.2759218215942383,
3.2743494510650635,
3.2755558490753174,
3.2739055156707764,
]
times = [
110410.41018200485,
110524.66757500588,
110523.31244999732,
110481.97357999743,
110394.39703700191,
]
print("p=%.4f" % scipy.stats.ttest_1samp(losses, 3.28, alternative="less").pvalue)
# p=0.0012
print("losses:", torch.std_mean(torch.tensor(losses)))
# losses: (tensor(0.0015), tensor(3.2755))
print("time:", torch.std_mean(torch.tensor(times)))
# time: (tensor(61.6309), tensor(110466.9531))My new machine took 20 more steps than reported by the prior WR, which is rather large, but it was consistent over all runs so I'll attribute it to inter-machine variance. Hence I updated the new steps to |
@ClassicLarry My intuition is that these gates function slightly better when their inputs are more "spread out", i.e. they have the chance to learn both shared and independent features. If there weren't enough gates, I believe you that moving the I didn't probably ablate what the "ideal overlap" is of shared/independent features, albeit I don't think there's much gains left in optimizing this. |
|
I’ll do some validation runs and ablations later. It’s not yet clear this is a new record since the step count is increased from 1775 and the loss is higher. Different machines shouldn’t cause a step count change, just a change in runtime. |
|
Fair point, I overrode the default harness to make running repeats easier -- I ran repeats as a for-loop over the Python train script, reinitializing the model every time. It's possible that changed the baseline architecture for non-obvious reasons. Thanks for checking, I'll try retiming on a fresh machine on my end too. |
|
Some runs: I dont think the gate update is doing anything. Its easy to see a drop from 3.28 to 3.276 for a single run and mistake noise for an improvement. The value embed update looks to be worthwhile, and makes the model arch more symmetric and simpler. I'll merge once you can show p value for the updated step count, for any combination of updates that works. For repeat runs you can do something like Not looking to merge in the harness code. If this doesn't quite get p value at 15 steps (I think it will), you can probably add mantissa to Adam or turn more short windows into PairedHeadLayers to get over the threshold. |
* PI config; Trying residual dims separation * Profiler * value embed is symmetric * Agg treatment before flight, ~ -20steps * Easy trials * add WR record; better trial * Cleanup for push * More cleanup * more --------- Co-authored-by: Ubuntu <[email protected]> Co-authored-by: Dummy Name <[email protected]> Co-authored-by: dummy user <email>
|
@ClassicLarry Thanks for timing my ideas. Unfortunately, after the other WRs in between these 2 weeks, this idea isn't significant anymore. I did further testing and found a new improvement. This is superceded by #209 |
2 participants
Updates on PR#190:
I can retime for PR#191 later.
Update: Retiming on PR#191 holds up on 15 steps and -0.7s improvement exactly.
Gates
Currently, we have$4$ very distinct gates: smear, skip, attention, and value embedding. For initial embedding $x_t$ at token $t$ , a smear gate shifts $x_{t-1}$ , and assigns $x_t := x_t + g \times x_{t-1}$ . The smear is a cheap proxy for the token shift that attention heads have already been observed to learn at higher cost. A skip gate enables a standard skip connection, and an attention gate controls the standard attention before MLP. Finally, a value embedding adds extra sparse embeddings that are mixed into the value of attention layers.
Each gate serves distinct roles; by the induction heads hypothesis, separate concerns live on separate subspaces of the residual vector. However, the WR had designed the gate to be a function of the first$c = 12$ dimensions, namely
$$g(x) = \sigma(w_g \cdot x_{1:c}), \quad w_g: \mathbb{R}^c$$ $x$ (with minor overlap).
which can "crowd" the residual stream with noise. To cleanly separate concerns, we modify gates to operate on distinct chunks of
If gates are "similar" enough in function, we find that some dim overlap can boost the learning signal, especially in early training. We arrived on this configuration below. Although we could grid search the ideal overlap, the main idea is to reduce the pressure on the first 12 dimensions in the residual stream, so this is a reasonable starting point.
Value Embeddings
After pruning some layers in prior WRs, the placement of the value embeddings became$W_v \in \mathbb{R}^{v \times d}$ , for embedding size $d = 768$ and vocabulary $v = 50257$ . These embeddings add token-specific information to enrich the block output. We applied this method in a U-Net style over the $11$ Transformer blocks, namely
where shared indices share$W_v$ . However, we noticed the 0th embedding was only used once, wasting a very large matrix, so we modified the pattern to be symmetric like
Each embedding matrix is$v \times d = 50257 \times 768 \approx 38.6$ million parameters, and we want to ensure huge parameters are used efficiently.
Timing
Retiming record of PR#190 at 1840 steps
New record timing at 1825 steps
If no changes, will merge at latest WR - improvement = 109.2s-(0.7s) = 108.5s to be consistent with 0.7s improvement.