Consider using Muon (PR seems ready, or stuck?):
or (I didn’t know of this, newer/better variant?):
You currently use Adam, why not AdamW (or Muon etc)?
I suppose you’re using SwiGLU (SwiGLU-Clip was new to me and), newer to consider:
I suppose your problem is the 512 token context window, one reason and/or small/bad training dataset(?), for not being better. I think it might be intentional because of quadratic attention. Too costly for you to increase? I suppose SubQ is too hard to implement. You can have some linear layers, if you don’t already (and/or Mamba layers).
This paper seems awesome, but I guess a premature to add to your design:
https://papers.ssrn.com/sol3/papers.cfm?abstract_id=7170898
We propose Hierarchical Adaptive Memory (HAM), a memory layer that tiers information into working, episodic, and semantic stores by a utility score (frequency, reuse, recency, novelty, predictive utility, stability) and consolidates frequent, redundant content into fewer prototypes while retaining rare content at higher fidelity. This predictive-coding-inspired compression lowers byte footprint at a similar fine-tuning budget. Unlike prior tokenization-level hierarchical memory, HAM operates at the prototype level and so can serve as an external store, a KV-cache compressor, or an internal memory block.
AVQ-Attention: Adaptive Vector-Quantized Attention
https://arxiv.org/pdf/2607.12789v1
seems like a very important paper (for you/LLMs too? It’s validated on image classification, AI suggested to me it would work for LLMs, and I challenged that, since not mentioned, only LLMs in sources, it thinks it translates to LLMs). If the I understand correctly and the codebook, M, is small, then much better than the quadratic attention I suppose you use, i.e. it’s O(MN).
CommVQ paper is also intriguing, maybe it can combine with above:
.. FP16 KV cache size by 87.5% with 2-bit quantization, while outperforming state-of-the-art KV cache quantization methods. Notably, it enables 1-bit KV cache quantization with minimal accuracy loss, allowing a LLaMA-3.1 8B model to run with a 128K context length on a single RTX 4090 GPU. The source code is available at: GitHub - UMass-Embodied-AGI/CommVQ: [ICML 2025] CommVQ: Commutative Vector Quantization for KV Cache Compression · GitHub.
In case you or anyone wants to do 2-bit quantized models (a bit premature until your generated model is better):
Experiments across diverse LLMs demonstrate that LC-QAT consistently outper-
forms state-of-the-art QAT methods while using only 0.1%–10% of the training data.
They use AdamW in the paper (I suppose your Adam would also work, and probably Muon or similar too).
Note, the paper was rejected (it seems very intriguing to me, found it reference in breakthrough paper below, your call if you want to implement it), but read arguments and answers, still starting this way in the metareview:
I agree with the reviewers that the idea of the paper is simple, elegant, and appealing. Adding a latent variable to the decoder with very small overhead is attractive, and the paper shows improvements on several benchmarks. In this sense, the work is a tangible contribution and shows promising results.
Explorative Modeling: Unlocking a Third Pretraining Axis and End-to-End Generation
Autoregressive LLMs have proven the hardest case, for the reasons discussed in our limitations (Section 7)
[..]
The Free Transformer [68] conditions a decoder on a latent variable inferred by a VAE, which is exactly the kind of latent exploration searches over, and training it with Explorative Modeling instead would remove the VAE entirely, along with the exposure bias of training on inferred latents (Section 3.1). Beyond language models, few-step models such as MeanFlow [15] are a natural fit for exploration as well, since exploration can supply the generative expressivity their shortened trajectories amortize. We also believe in combining XMs with Energy-Based Transformers (EBTs) [69], where the biggest documented challenge with EBTs has been end-to-end generation and handling highly multimodal distributions, which is exactly what XMs enable. Paired together, XMs and EBTs could enable more dynamic reasoning, search, and generalization over entire sequences.
[Note e.g. there Figure 9 needs light mode, broke in dark mode]
The longer a model generates, the worse things get—you’ve probably seen this firsthand with video generation models that melt into mush after ten seconds, or LLMs that get less coherent over really long generations.
TLDR: We outscale (feed-forward) transformers while generalizing reasoning/system 2 thinking to any modality/problem without requiring verifiable rewards
. Energy-Based Transformers are the first approach to outscale feed-forward transformers across modalities and with respect to several axes including data, depth, parameters, FLOPs, etc. Energy-Based Transformers can think over every single prediction (i.e. every token in language modeling) and generalize better than existing models.
https://arxiv.org/pdf/2507.02092
Specifically, we train Energy-Based Transformers (EBTs)—a new class of Energy-Based Models (EBMs)—to assign an energy (un-
normalized probability) value to every input and candidate-prediction pair, enabling
predictions through gradient descent-based energy minimization until convergence.
This formulation enables System 2 Thinking to emerge from unsupervised learn-
ing, making it modality and problem agnostic. Across both discrete (text) and
continuous (visual) modalities, we find EBTs scale faster than the dominant Trans-
former++ approach during training, achieving an up to 35% higher scaling rate with
respect to data, batch size, parameters, FLOPs, and depth. During inference, EBTs
improve performance with System 2 Thinking (i.e., extra computation) by 29%
more than the Transformer++ on language tasks, and EBTs outperform Diffusion
Transformers on image denoising while using fewer forward passes