I don’t think we need to limit to pure Julia, could e.g. reuse CUDA/Triton/PyTorch kernels (often super-optimized, and well already existing…) etc. for just catching up. At least for optional parts.
@mantzaris now that we have your LLM code from scratch in Julia, we’ve been discussing improvements, I just hesitate to post a lot of (more) stuff in the package announcement thread itself.
I’m looking for breakthroughs that do not mean we will be playing catch-up, or can avoid implementing some code/model arch that may already be outdated. One possibility is implementing FlashAttention and/or FlexAttention, but maybe defer, in case attention itself is changing a lot. The holy grail I think is memory, as in learning over time, “continual learning”, in some form and one paper related to that (https://www.youtube.com/watch?v=P9uNy71YukQ&t=320s):
However, our model continues learning at test time via next-token prediction on the given context, compressing the context it reads into its weights.
Regarding attention, maybe implement TANGO and/or WANGO:
A standard Transformer block separates cross-token interaction in self-attention from the nonlinear feed-forward network applied independently at each position. We introduce Token-Aggregated Nonlinear Gating Operators (TANGO), which replaces these two sublayers with a single cross-token gated residual update. Each source token produces a Swish-gated linear unit (SwiGLU) gate vector.
..
The Tango model computes a separate weight for every causally visible source. Windowed Aggregation of Nonlinear Gating Operators (WANGO) retains the same unnormalized pairwise scores within a recent window. For older sources, it uses a positive feature-map weighting rule whose required statistics can be maintained with running sums. Tango is quadratic in sequence length. For fixed window and feature dimensions, the Wango model is linear in sequence length.
..
We compare TANGO and WANGO with Recurrent and Untied Transformer++, full-attention GAU, and FLASH. All models have approximately 44.3M nonembedding parameters and are trained in three matched runs. TANGO, WANGO, and Recurrent Transformer++ apply one shared block four times; the other architectures use four independent blocks.
..
On FineWeb-Edu, the Wango model obtains the lowest mean validation negative log-likelihood (NLL) among models whose computation is linear in sequence length. It also has lower mean NLL than Recurrent Transformer++ at nearly the same analytical forward-pass multiply–accumulate count. The Tango model obtains the lowest mean validation NLL on FineWeb-Edu, Lean, and DeepMind Mathematics, but has the largest analytical forward-pass operation count among the compared architectures
All models use context length 256. Training lasts 84,000 optimizer steps with an effective batch size of 64 examples, corresponding to 5.376 million examples or 32,000 training examples per module–difficulty combination.
This small context might be helpful to you, but note there it applies to DeepMind Mathematics, and I see up to FineWeb-Edu (8,192-token context; 5,722 steps) in figure 2. Figure 5 is also interesting.
“DeepMind Mathematics” is apparently this dataset: GitHub - google-deepmind/mathematics_dataset: This dataset code generates mathematical question and answer pairs, from a range of question types at roughly school-level difficulty. · GitHub
Even if you don’t change your code at all, then it might be good to know of and use such datasets for you.
above paper references this one:
I’ve been very intrigued by looped models, and this seems like an improvement (I added bold):
Importantly, T2MLR does not require pretraining from scratch: retrofitting the recurrent pathway into an existing pretrained 1.7B Transformer and briefly finetuning substantially improves math reasoning, lowering the barrier to practical adoption. These results suggest that effective latent reasoning in Transformers does not require looping over all layers as in previous works, but can instead emerge more strongly from targeted middle-layer recurrence.