Techniques

Every non-trivial technique used in Lavender-2 v0.1.0, with a pointer to the paper that introduced it and the reference implementation the code mirrors. If a future agent wants to change one of these, they should read the primary source first.

All links below are to public artifacts. Paper links point to arXiv where available.

---

Tokenizer

Byte-Pair Encoding with a byte-level pre-tokenizer

• What: BPE vocabulary learned from scratch on the normalized training corpus. Byte-level pre-tokenization (ByteLevel) makes the tokenizer lossless over arbitrary UTF-8 input.
• Paper: Sennrich, Haddow & Birch, Neural Machine Translation of Rare Words with Subword Units, ACL 2016. https://arxiv.org/abs/1508.07909
• Byte-level variant: Radford et al., Language Models are Unsupervised Multitask Learners (GPT-2 tech report), 2019. https://cdn.openai.com/better-language-models/language_models_are_unsupervised_multitask_learners.pdf
• Reference code: HuggingFace tokenizers library, https://github.com/huggingface/tokenizers. Lavender-2 uses models.BPE, pre_tokenizers.ByteLevel, decoders.ByteLevel, processors.ByteLevel, and trainers.BpeTrainer directly.

---

Architecture

Decoder-only Transformer

• Paper: Vaswani et al., Attention Is All You Need, NeurIPS 2017. https://arxiv.org/abs/1706.03762
• Decoder-only variant / LM framing: Radford et al., Improving Language Understanding by Generative Pre-Training (GPT-1), 2018. https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf
• Reference code: HuggingFace transformers/models/gpt2 and models/llama (the closer structural analogue: RMSNorm + RoPE + SwiGLU).

RMSNorm (pre-norm)

• What: Replaces LayerNorm with a root-mean-square rescaling; no mean centering and no learnable bias. Applied pre-block (Xiong et al. pre-norm recipe) rather than post-block.
• Paper: Zhang & Sennrich, Root Mean Square Layer Normalization, NeurIPS 2019. https://arxiv.org/abs/1910.07467
• Pre-norm positioning: Xiong et al., On Layer Normalization in the Transformer Architecture, ICML 2020. https://arxiv.org/abs/2002.04745
• Reference code: Meta's Llama, RMSNorm in https://github.com/meta-llama/llama/blob/main/llama/model.py.

Rotary Position Embeddings (RoPE)

• What: Apply a position-dependent rotation to query and key vectors in each attention head; no additive positional embedding on the input. Lavender-2 precomputes cos / sin buffers once and applies them per layer.
• Paper: Su et al., RoFormer: Enhanced Transformer with Rotary Position Embedding, 2021. https://arxiv.org/abs/2104.09864
• Reference code: apply_rotary_emb in Llama, https://github.com/meta-llama/llama/blob/main/llama/model.py.

Grouped-Query Attention (GQA)

• What: Share KV heads across groups of query heads (n_heads queries, n_kv_heads < n_heads keys/values). Cuts KV-cache memory at inference with minimal quality loss versus multi-head attention.
• Paper: Ainslie et al., GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints, EMNLP 2023. https://arxiv.org/abs/2305.13245
• Precursor (MQA): Shazeer, Fast Transformer Decoding: One Write-Head is All You Need, 2019. https://arxiv.org/abs/1911.02150
• Reference code: Llama 2 attention block, https://github.com/meta-llama/llama/blob/main/llama/model.py.

KV cache (inference)

• What: Cache K and V per layer across autoregressive steps so each new token costs one forward pass at the new timestep only.
• Standard autoregressive-decoding trick; first appears as a practical optimization in OpenAI / Google codebases rather than a named paper. The HuggingFace documentation has a clean writeup: https://huggingface.co/docs/transformers/main/kv_cache.

SwiGLU feed-forward

• What: FFN of the form down_proj(silu(gate_proj(x)) * up_proj(x)). Replaces the standard Linear -> GELU -> Linear MLP.
• Paper: Shazeer, GLU Variants Improve Transformer, 2020. https://arxiv.org/abs/2002.05202
• Reference code: Llama FeedForward, https://github.com/meta-llama/llama/blob/main/llama/model.py.

Mixture of Experts with top-k routing

• What: N SwiGLU experts per FFN slot; a learned linear gate softmax-routes each token to its top-k experts. Outputs are combined by the renormalized gate weights. A load-balancing auxiliary loss penalizes collapsed routing distributions.
• Switch Transformer (top-1 routing, aux loss we use): Fedus et al., Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity, 2021. https://arxiv.org/abs/2101.03961
• Sparsely-Gated MoE (original top-k recipe): Shazeer et al., Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer, ICLR 2017. https://arxiv.org/abs/1701.06538
• Top-2 at scale: Jiang et al., Mixtral of Experts, 2024. https://arxiv.org/abs/2401.04088
• Reference code: Google's Mesh-TensorFlow MoE, https://github.com/tensorflow/mesh, and Mistral's Mixtral reference, https://github.com/mistralai/mistral-inference.

---

Training

AdamW optimizer

• Paper: Loshchilov & Hutter, Decoupled Weight Decay Regularization, ICLR 2019. https://arxiv.org/abs/1711.05101
• Reference code: torch.optim.AdamW.

Cosine-annealing learning rate

• Paper: Loshchilov & Hutter, SGDR: Stochastic Gradient Descent with Warm Restarts, ICLR 2017. https://arxiv.org/abs/1608.03983
• Reference code: torch.optim.lr_scheduler.CosineAnnealingLR.

bfloat16 mixed-precision training (AMP)

• What: Forward + loss in bf16 via torch.autocast; optimizer state stays in fp32; GradScaler stabilizes the backward pass when any fp16 op is still in play.
• Paper (bf16 rationale): Kalamkar et al., A Study of BFLOAT16 for Deep Learning Training, 2019. https://arxiv.org/abs/1905.12322
• AMP recipe paper: Micikevicius et al., Mixed Precision Training, ICLR 2018. https://arxiv.org/abs/1710.03740
• Reference code: torch.amp.autocast, torch.amp.GradScaler.

Distributed sharding (FSDP / DeepSpeed ZeRO-3)

Target runtime for the 70B config; not exercised by the dev config.

• ZeRO: Rajbhandari et al., ZeRO: Memory Optimizations Toward Training Trillion Parameter Models, SC 2020. https://arxiv.org/abs/1910.02054
• FSDP: Zhao et al., PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel, VLDB 2023. https://arxiv.org/abs/2304.11277
• Reference code: torch.distributed.fsdp.FullyShardedDataParallel (https://pytorch.org/docs/stable/fsdp.html) and DeepSpeed ZeRO-3 (https://github.com/microsoft/DeepSpeed).

---

Datasets (v0.1.0 mix)

Every dataset in the registry, with its source and citation. Access rules differ: gated datasets require a HuggingFace token and manual terms acceptance.

Key	HF path	Paper / source	Access
wildchat	allenai/WildChat-1M	Zhao et al., WildChat: 1M ChatGPT Interaction Logs in the Wild, ICLR 2024. https://arxiv.org/abs/2405.01470	open
soda	allenai/soda	Kim et al., SODA: Million-scale Dialogue Distillation with Social Commonsense Contextualization, EMNLP 2023. https://arxiv.org/abs/2212.10465	open
dailydialog	roskoN/dailydialog (parquet mirror)	Li et al., DailyDialog: A Manually Labelled Multi-turn Dialogue Dataset, IJCNLP 2017. https://arxiv.org/abs/1710.03957	open
topical_chat	agentlans/Conversational-Reasoning-Topical-Chat	Gopalakrishnan et al., Topical-Chat: Towards Knowledge-Grounded Open-Domain Conversations, Interspeech 2019. https://www.isca-archive.org/interspeech_2019/gopalakrishnan19_interspeech.html	open
multiwoz	pfb30/multi_woz_v22	Zang et al., MultiWOZ 2.2: A Dialogue Dataset with Additional Annotation Corrections and State Tracking Baselines, NLP4ConvAI 2020. https://arxiv.org/abs/2007.12720	open
taskmaster	google-research-datasets/taskmaster2	Byrne et al., Taskmaster-1: Toward a Realistic and Diverse Dialog Dataset, EMNLP 2019. https://arxiv.org/abs/1909.05358 (v2 is the dataset shipped). https://github.com/google-research-datasets/Taskmaster	open
wikipedia	wikimedia/wikipedia (20231101.en)	Wikimedia Foundation, plain-text English Wikipedia dump. https://dumps.wikimedia.org/	open
lmsys_chat	lmsys/lmsys-chat-1m	Zheng et al., LMSYS-Chat-1M: A Large-Scale Real-World LLM Conversation Dataset, ICLR 2024. https://arxiv.org/abs/2309.11998	gated
arena	lmsys/chatbot_arena_conversations	Zheng et al., Judging LLM-as-a-Judge with MT-Bench and Chatbot Arena, NeurIPS 2023. https://arxiv.org/abs/2306.05685	gated
claude_opus	Roman1111111/claude-opus-4.6-10000x	Community distillation set of Claude Opus 4.6 reasoning traces (math & logic); no formal paper.	open
kimi_k25	ianncity/KIMI-K2.5-700000x	Community distillation set of Moonshot AI's Kimi K2.5 reasoning traces; no formal paper.	open
opus_reasoning	nohurry/Opus-4.6-Reasoning-3000x-filtered	Filtered community distillation of Claude Opus 4.6 reasoning with thinking traces; no formal paper.	open

Dataset weights (per-dataset weight in the notebook registry) determine the mixture proportions at pretrain time. See the registry in full_llm_pipeline.ipynb (search for "wildchat").

---

Notes on provenance

The "Reference code" links above point to the implementation the Lavender-2 notebook code most closely mirrors. Where Meta's Llama repository is listed, the specific functions referenced (RMSNorm, apply_rotary_emb, FeedForward, attention block) are the canonical readable implementations — the notebook code is written independently in plain PyTorch rather than copied, but the shape and variable names track Llama conventions deliberately so a reader coming from that codebase doesn't have to retranslate.