To read

Attention

MHA

Attention is dense:

prefill cost: \(O(n^2d)\)

KV cache per token: \(O(2hd_h) = O(2d)\)

best baseline quality but long-context inference bottlenecked by dense attention and large KV cache

MQA

reduce KV cache by sharing keys/values across head

• MQA: all query heads share one KV head
• GQA: query heads are grouped. each group shares one KV head

prefill cost: \(O(n^2d)\)

KV cache per token: \(O(2h_{kv}d_h)\), so MQA/GQA mostly reduce decode-time memory bandwidth

MLA

(DeepSeek-V2)

For head $i$, each key is split/concatenated as $k_{j,i} = [k_{j,i}^C; k_j^R]$. Attention score is:

MLA project tokens to compressed latent state: \(c_t^{KV} = x_t W^{DKV}\) then, reconstruct keys/values from latent state:

so during inference we cache \(c_t^{KV}\) instead of all \((k_{t,i}, v_{t,i})_{u=1}^h\)

prefill cost: \(O(n^2d)\)

KV cache per token: \(O(d_c)\), so even further reduce memory bandwidth

NSA (Native Sparse Attention)

(Native Sparse Attention)

For each token $t$, NSA computes three attention branches and mixes them:

where each gate $g_t$ is learned from token representation.

Uses 3 branches: compressed global tokens, selected fine-grained blocks, sliding-window local tokens

• Compressed Attention

  Compress past keys/values into overlapping blocks:

  \[\tilde{k_t}^\text{cmp} = \phi_K(k_{id+1:id+l}), \qquad \tilde{v_t}^\text{cmp} = \phi_V(v_{id+1:id+l})\]

  with $0 \le i \le \lfloor \frac{t-l}{d}\rfloor$. Usually $l > d$ because we want overlap.

  Attend over compressed block summaries:

  \[o_t^\text{cmp} = \text{Attn}(q_t, \tilde{K}_t^\text{cmp}, \tilde{V}_t^\text{cmp})\]

  which gives cheap global context.

  Time Complexity: $O(\frac{n^2}{d}d_h)$

• Selected Attention Use compressed attention scores to decide which original token blocks matter:
  \[p_t^\text{cmp} = \text{softmax}\Bigg(\frac{q_t^\intercal \tilde{K}_t^\text{cmp}}{\sqrt{d}}\Bigg)\]

  Convert these into scores for original token blocks, choosing the top $n_\text{slc}$ blocks:

  \[I_t = \{i:\text{rank}(p_t^\text{slc}[i]) \le n_\text{slc}\}\]

  This comes to around $n_\text{slc}l$ uncompressed tokens in which:

  \[o_t^\text{slc} = \text{Attn}(q_t, \tilde{K}_t^\text{slc}, \tilde{V}_t^\text{slc})\]

  Time Complexity: $O(n n_\text{slc} l d_h)$

• Sliding Window Attention

  Always attend to recent local tokens:

  \[\tilde{K}_t^\text{win} = K_{t-w:t}, \qquad \tilde{V}_t^\text{win} = V_{t-w:t}\] \[o_t^\text{win} = \text{Attn}(q_t, \tilde{K}_t^\text{win}, \tilde{V}_t^\text{win})\]

  Time Complexity: $O(nwd_h)$

So total time complexity is:

DSA (deepseek sparse attention)

(DeepSeek-V3.2)

DSA has two additional parts: lightning indexer and top-k sparse attention.

• Lightning Indexer: for query token $t$, score each previous token $s$: \[I_{t,s} = \sum_{j=1}^{H_I} w_{t,j}^I \text{ReLU}(q_{t,j}^I \cdot k_s^I)\]

  where $H_I$ is the number of indexer heads. Each indexer retrieves on different criteria.

  Here we get $q_{t,j}^I, w_{t,j}^I$ from the current query token and $k_s^I$ comes from each previous token. Use ReLU for throughput as indexer is lightweight enough to run efficiently in FP8.

  Time Complexity: $O(n^2H_Id_I)$

• Top-k Token Selection: using the index scores we select the most relevant previous tokens
  \[S_t = \text{TopK}_s(I_{t,s},k)\]

  We only retrieve those KV entries:

  \[K_t^\text{DSA} = K_{S_t}, \qquad V_t^\text{DSA} = V_{S_t}\]

  So normal attention is computed only over the selected tokens:

  \[o_t = \text{Attn}(q_t, K_t^\text{DSA}, V_t^\text{DSA})\]

  Time Complexity: $O(nkd)$

So total time complexity is

DeepSeek-V3.2 instantiates MLA under MQA mode.

• Sparse training uses selected tokens and trains whole model:

  \[\mathcal{L} = \mathcal{L}_{LM} + \mathcal{L}_\text{indexer}\]

• select 2048 KV tokens per query in sparse training

Architecture

Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention

What's the problem

• Vanilla Transformer self-attention is quadratic in sequence length during training/inference
• Autoregressive inference gets painful as each token makes KV grow linearly

Goal: causal (autoregressive attention) in linear time

Core Idea

replace softmax attention with a kernel feature map →attention is associative, computing output using a running state instead of NxN attention matrix

If we instead use \(\text{sim}(q,k) = \phi(q)^\intercal \phi(k)\), with feature map \(\phi\), attention can be rewriting as

During autoregressive attention, we define

• \(S_i = \sum_{j \le i}\phi(K_j)V_j^\intercal\)
• \(Z_i = \sum_{j \le i}\phi(K_j)\)

Then the output is:

and we can update our states:

• \(S_i = S_{i-1} + \phi(K_i)V_i^\intercal\)
• \(Z_i = Z_{i-1} + \phi(K_i)\)

this is a constant size in memory, unlike the linear growth in regular attention.

Training: still parallelizable across positions as prefix sums are computed efficiently

Inference: per-token cost becomes constant wrt to context length

Massive speedups for long autoregressive sequences while attaining similar performance to vanilla Transformers

Linear Transformers Are Secretly Fast Weight Programmers

What's the problem

Linear attention gives O(n) time and can be run as an RNN-style recurrence, but in practice it often underperforms softmax attention on harder long-context behaviors.

Paper's idea: Argues that a core reason why underperformance occurs is how memory is written. Standard linear-attention update is purely additive, making it hard to edit/correct stored key→value associations once capacity is stressed

Core Idea

causal linear Transformer is a fast-weight memory system: it maintains matrix state updated by outer products (keys \(\otimes\) values).

• view it as “fast weights” exposes capacity/overwriting issue
• Better update rule: delta-rule-style with a learned learning rate so the model can revise memory instead of only accumulate

  Linear attention = “programming instructions” on a memory matrix

• maintain a matrix memory \(W_t\), plus a normalizer state
  • aka S and Z
• each time step write an outer product of a key and a value into that memory

Fast-weight view, a token emits instructions:

• “write” instruction: add \(\phi(k_t)v_t^\intercal\) to \(W\)
• “read” instruction: query with \(\phi(q_t)^\intercal W\) (plus normalization)

→ linearised Transformers are mathematical identical to classic fast-weight systems

Fast weight equivalence

Interpret \(S_t\) as a fast weight matrix \(W_t\).

Read: \(r_t = \tilde{q}_t^\intercal W_t\), output uses normalization \(y_t = \frac{r_t}{\tilde{q}_t^\intercal z_t}\)

Write: \(W_t = W_{t-1} + \tilde{k}_t v_t^\intercal\)

Problem: Additive writes + finite memory $\rightarrow$ overcapacity

• \(W_t\) has fixed size, so as sequence length grows, end up with many keys competing inside the same matrix
• “wrong”/stale associations will linger with only additive updates, which results in interference
• model should be able to keep/delete/correct dynamically

Fix: Delta-rule programming + learned per-step learning rate

• instead of always add \(kv^\intercal\), write an error-correcting instruction that can move the current mapping to desired association
• model should learn a dynamic learning rate that changes over time, so it can decide when to aggressively overwrite or gently integrate

Let the current memory have prediction for value given key:

Define error \(e_t = v_t - \hat{v}_t\). Update fast weights with learned step size \(\eta_t\) with

Delta Network (DeltaNet)

• standard linear-attention normalization can be added to delta fast weights, but it causes instability and breaks write/remove balance
• normalize keys and queries to sum to 1
• preserves attention interpretation while keeping delta rule well-behaved

Proposed Alternative Feature Map to FAVOR+

FAVOR+:

Gated Linear Attention Transformers with Hardware-Efficient Training

What's the problem

vanilla linear attention often lags softmax on langauge modeling (in quality). Also, many linear-attention implementations aren't I/O-aware, so they can be slower than optimized softmax attention.

Paper claim to fix both with:

• FlashAttention-style kernel for linear attnetion
• data-dependent gating parameterization that boosts expressivity while staying trainable efficiently

Core Ideas

• FLASHLINEARATTENTION: makes linear attention really fast on GPU by tiling + SRAM reuse + chunkwise recurrence/parallelism tradeoffs
• GLA (Gated Linear Attention): upgrades linear-attention recurrence with a data-dependent 2D forget gate so the model can selectively retain/forget parts of the matrix state (like an LSTM, but in “fast-weights” state)

Linear Attention as a 2D-state RNN

• linearize softmax attention via kernel feature space \(\phi\), turning attention into something with a matrix state.

Gated Linear Attention (GLA)

introduce a 2D forget gate \(G_t \in (0,1)^{d_k \times d_v}\) so recurrence is now:

So instead of always accumulating into \(S\), you can decay/retain specific subentries of the state in a **data-dependent way

• more expressive than using single global decay

FlashLinearAttention

• training bottleneck is HBM ↔ SRAM traffic and whether you can use tensor cores
• Chunkwise-algorithm
  • split sequence into chunks of size C
  • Training is
    • interchunk: carry a chunk-level state forward (RNN-like)
    • intrachunk: do dense matmul inside each chunk (attention-like), with causal mask

write chunk output as:

and complexity becomes \(O(LCd + Ld^2)\), which beats \(O(L^2d)\)

With chunkwise matmuls, you pick C to be a multiple of 16 to hit tensor cores. There's a trade off

• more parallelism (larger C, more like attention)
• less memory (smaller C, more like an RNN)

Two kernel variant

• Non-materialization: keep the running S in SRAM, and stream chunks sequentially → memory efficient, but less sequence-parallel
• Materialization: first write chunk states \(S[n]\) to HBM (so chunks can be computed independently), then compute chunk outputs in parallel across chunks → better occupancy when batch is small

Showed:

• FLA kernel is faster than FA-2 as a standalone layer even at \(\sim\)1K length
• for moderate-scale LM, GLA is competitive vs strong LLaMA-style Transformer baseline and vs linear-time baselines like RetNet and Mamba
• Length generalization: report a model trained on 2K generalizing to \textgreater20K without big perplexity degradation
• Throughput is higher than similarly sized Mamba

Hydra: Bidirectional State Space Models Through Generalized Matrix Mixers

What's the problem

SSMs like Mamba are inherently causal, which is awkward for non-causal tasks (masked LM, classification, vision) where bidirectional context matters

• want to use a structured-matrix lens to design a principled bidirectional extension that keeps SSM efficiency and improves expressivity

Core idea: view sequence mixers as an \(L \times L\) mixer matrix \(M\) acting on a projected sequence, then choose a structured matrix class that is both bidirectional and efficient.

• Hydra chooses quasiseparable matrices, a generalization of both low-rank (linear attention) and semiseparable (SSMs)
  • all strictly lower- and upper-triangular submatrices have low rank, independent of T

Matrix Mixer Framework

A sequence mixer:

where \(X \in \mathbb{R}^{L \times C}\) (L elements in input sequence, each with C channels), \(f_X\) is preprocessing/projections, and each head has a mixer matrix \(M^h \in \mathbb{R}^{L \times L}\).

• They claim that adding SAM structure materially boosts performance

State Space Models are Semiseparable Matrix Mixers

Recap the SSM-as-matrix view:

with elements

This implies M is (lower-triangular) semiseparable

• semiseparable (SS): any submatrix taken from the lower-triangle (including diagonal) has rank \(\le N\)

Limitation is that this structure bakes causality

Hydra's structured-matrix: quasiseparable = principled bidirectionality

They choose QS matrices because

• include upper + lower structure (bidirectional)
• expressive
• have sub-quadratic algorithms (instantiations are linear-ish)

actual QS parameterization

define an N-quasiseparable matrix by speecifying each entry:

Quasiseparable (QS): any submatrix from the strictly upper or strictly lower triangle has rank \(\le N\)

• both past→future and future→past interactions factor through a small state
• you can realize the operator with two scans (forward + backward) or an equivalent coupled state

This is the right generalization as

• QS generalizes low-rank mixers (linear attention)
• QS generalizes and extends semiseparable mixers (SSMs)

QS is strictly more expressive than common SSM heuristics as those constrain the diagonal value, whereas QS enables free diagonals.

• diagonal actually apparently explains performance gains

Taming the Hydra: how to implement QS efficiently

The action of a QS matrix can be decomposed into two SS applications plus simple sequence ops + a diagonal term:

where \(D = \text{diag}(\delta_1, \dots, \delta_L)\), flip reverses the sequence, and shift right-sfhits with zero pad

• first term approximately forward scan contributions
• second SS term on flipped sequence approximately is backward scan contributions
• DX gives the free diagonal terms

so Hydra can reuse a strong SSM kernel and just add flip/shift/diag glue

Note:

• parameter increase over single SSM is marginal as QS modeling encourages sharing projection

Gated Delta Networks: Improving Mamba2 with Delta Rule

Problem

Mamba2 and DeltaNet solve complementary parts of the linear-time sequence modeling problem.

Mamba2 uses a gated decay mechanism, so it can quickly forget stale context:

This is useful for context switching, but the forgetting is coarse: the same decay can remove useful information along with irrelevant information.

DeltaNet instead treats the recurrent state as an associative memory from keys to values. Given key \(k_t\), value \(v_t\), and write rate \(\beta_t\), the delta rule is

Equivalently,

The term \(S_{t-1}k_t\) is the current prediction for key \(k_t\), so the update only writes the residual error. This makes DeltaNet good at precise key-value editing, but without strong decay it can accumulate stale associations and struggle under context switches.

Core Idea

Gated DeltaNet combines both mechanisms: keep DeltaNet's targeted key-value update, but add a Mamba2-style decay gate on the old state:

Here

The delta term $\beta_t\bigl(v_t - S_{t-1}k_t\bigr)k_t^\top$ controls precise associative editing.

The readout is still linear-attention style $o_t = S_t q_t$ up to implementation details such as normalization, output gating, and projection.

Chunkwise Parallel Training

Token-by-token, the recurrence is sequential. To train efficiently, split the sequence into chunks of length \(C\). Let \(S[t]\) be the state entering a chunk, and let \(r=1,\dots,C\) index tokens inside the chunk.

Expanding the recurrence within one chunk gives

This separates the chunk state into two pieces:

Naively, this is still expensive because it contains many products of rank-one update matrices: $I-\beta_i k_i k_i^\top$.

The trick is to use a WY-style representation, which rewrites these products using chunk-level matrices such as

Then the whole chunk can be computed with triangular solves and matrix multiplications instead of a Python-level scan over tokens.

A compact view of the chunk computation is

where \(\Gamma[t]\) stores the cumulative decay ratios from the \(\alpha\)-gates. This triangular solve gives the effective within-chunk value updates after accounting for previous writes and decay.

Then the outputs for all tokens in the chunk can be computed as

where \(M\) is the causal lower-triangular mask. The first term reads from the state entering the chunk; the second term adds all causal within-chunk updates.

The state passed to the next chunk is

Thus, Gated DeltaNet behaves like a recurrent associative memory during decoding, but training can be done with chunkwise parallel linear algebra. This keeps linear-time inference while making training GPU-efficient.

Test-time regression: a unifying framework for designing sequence models with associative memory

[TODO]

Log Linear Attention

Problem

Softmax attention is $O = \text{softmax}(QK^\intercal)V$, which is good expressivity but

Linear attention/Mamba-style models:

which can be recurrently written which is efficient. However this compresses all past tokens into one fixed-size state: $S_t \in \mathbb{R}^{d_v \times d_k}$.

Core Idea

We replace one state with logarithmically many states:

where $L = O(\log T)$.

Each state summarizes a Fenwick-tree bucket of the past:

which then outputs:

Equivalently:

ex

So it gets multiple nonuniform history summaries and chooses among them at output time.

NOTE GLA has one continuous decayed state across time.

Matrix Form

The model still fits:

For basic log-linear attention:

with

So

Complexity

Training uses chunking, so

The cost is

for fixed chunk size $C$: $O(T \log T)$.

Results

Interp

NLA (Natural Language Autoencoder)

two LLM modules:

• activation verbalizer (AV) maps activation to text description
• activation reconstructor (AR) mapping description back to activation

train both AV and AR with RL to reconstruct residual stream activation

• resulting NLA explanations read as plausible interpretations of model internals (does it really?)

LM have high-dim activation vectors (but are numbers aka opaque to human reader)

• training process doesn't explicitly incentivize NLA explanations to be interpretable or faithful
  • as training across Haiku3.5, 4.5, and Opus 4.6 there is more informative signals

example of NLA:

• while NLA claims are false, they are still faithful to context

we can also surface unverbalized aspects of model cognition (suspecting it is being tested)

• surface evaluation awareness more often on evaluation transcripts than real deployment traffic

still NLAs are blackboxes (so there is lack of mechanistic grounding), cost is high, excessive expressivity (on top of an activation)

training mechanism:

Model Training

Composer 2

Base model selection:

• starts from Kimi K2.5 (1T params, 32B active, 256K context).
  • chosen from both infrastructure fit and initial benchmark scores
• MLA, 61 layers

• coding-focused pre-training stage builds domain knowledge. More tokens here translate into measurably higher rewards after the RL stage.
• Three Phases:
  • Standard short-context Pre-training
  • Long-context extension
  • Final SFT phase

• We want agents to work until to fully complete a hard task
• Reinforcement learning is used to simulate user queries in-domain
  • Allows us to tune behaviors of the model
• Composer 1.5 bootstraps each training environment by exploring the repo, generating install commands, and writing verification tests before RL begins.

Reward shaping and self-summarization:

• A nonlinear length penalty balances speed and depth, while self-summarization lets the model continue past its context limit and still share one final reward across the rollout.

Kimi K2

Pretraining

Parameters: $1.04$ trillion-parameter MoE with about $32$B activated parameters.

Sparsity Scaling Law: Specific for MoE, sparsity is defined as ratio of total number of experts to number of activated experts. Using constant FLOPs (aka fixed number of experts) but increasing number of experts, it consistently lowers both training and validation loss.

They pretrained on 15.5T tokens with no loss spike:

On data, they rephrase instead of repeating (over epochs).

Infra

Cluster of $H800$ GPUs, with each node containing $2$TB RAM and $8$ GPUs connected via NVLink and NVSwitch. Across nodes, $8 \times 400$ Gbps RoCE interconnects are used.

Per model,

EP

So $256$-GPU model parallel group with ZeRO-1 data parallelism, which is $256D$ where $D$ is data-parallel degree.

IMPORTANT: Pipeline Parallelism:

We split global batch into micro-batches (kinda like comp-arch pipelining), but global batch gradient is still accumulated over all micro-batches.

\[\text{micro-batch size}=\frac{\text{global batch size}}{\text{data-parallel size} \times M }\]

• first few steps are underutilized even with micro-batches as pipeline needs to fill aka warmup bubble.

  Bubble fraction $\sim \frac{P-1}{M+P-1}$ where $P$ is the number of pipeline stages.

  Hence we must have $M \gg P$ for a small bubble. However high $M$ also hurts kernel efficiency and increase communication overhead so there is a tradeoff.

• Use 1F1B instead of all forward batches then all backwards batches. All forward and all backward is known as GPipe, but despite it keeping GPUs busy, it has a major memory problem ($\propto M$).

  1F1B has stage alternate between forward for a newer micro-batch and backward for an older micro-batch. This reduces activation memory while keeping pipeline utilization high.

  
  NOTE Loss is still the same, as we only update at the end of the batch
• Virtual pipeline stages instead of physical. Each physical rank is split into smaller layer chunks.
  

  Forward path is more interleaved, which gives more scheduling granularity.

  \[S_1^{(a)}\rightarrow S_2^{(a)} \rightarrow S_3^{(a)}\rightarrow S_4^{(a)}\rightarrow S_1^{(b)} \rightarrow S_2^{(b)}\rightarrow S_3^{(b)}\rightarrow S_4^{(b)}\]

Post-training

Post-training stack:

• SFT on normal instruction + agentic tool-use data
• Synthetic tool-use trajectory generation
• RL over verifiable and non-verifiable tasks

How they use agents for generating SFT tool-use data:

RL Objective:

difference between the two

GRPO uses the probability ratio directly. K2 regresses the log probability ratio to the reward advantage

\[\tau \log \frac{\pi_\theta(y_i \vert x)}{\pi_\text{old}(y_i \vert x)} \approx A_i\]

log is important since advantage $A_i$ can be positive or negative.

This tells us that psoitive advantage should add to token/sequence log-prob, and engative advantage should subtract from it.

For tasks that aren't verifiable they use a critic version of K2 to rank outputs with rubrics.

MAI-Thinking-1

[TODO]

SWE-1.7

SWE-1.7

Top-K sampling: In RL

where $e_j$ is one-hot vector for sampled token.

For a sampled token $j$:

• positive advantage $\tilde{A} > 0$: increase sampled token, decrease others
• negative advantage $\tilde{A} < 0$: decrease sampled token, increase others

So if a tail token is sampled and gets low reward, it not only penalizes the bad tail sample but boosts the already-dominant token most, which worsers entropy.

ex

With top-$p$ sampling, we exclude the tail which reduces mode amplification.

SSL

A Simple Framework for Contrastive Learning of Visual Representations

Core modeling recipe (SimCLR)

• two branches, shared encoder: take image → apply two random augmentations → encode both views
• add a small projection head (MLP) on top of the encoder during pretraining; contrastive loss is applied in that space
• uses NT-Xent / InfoNCE-style contrastive loss

Findings

• Augmentations are the “task definition”, performance is highly sensitive to augmentation pipeline
• With right augmentation + training setup, contrastive SSL yields representations that transfer well to downstream vision tasks under linear evaluation / finetune protocols

Data used:

• linear probing on ImageNet-style settings

InfoNCE

NT-Xent

Let minibatch have N images. Apply 2 random augmentations to each image → 2N views

Encode each view and project:

where f is encoder, g is projection MLP, and z is L2-normalized

Similarity is defined by cosine similarity

For a positive pair, per-example NT-Xent loss is:

Per batch loss,

assuming every two views comes from the same image

Momentum Contrast for Unsupervised Visual Representation Learning

Core Modeling recipe (MoCo)

• Two encoders with same architecture:
  • query encoder \(f_q(\cdot;\theta_q)\)
  • key encoder \(f_k(\cdot;\theta_k)\) (updated by momentum / EMA)
• For each image, make two augmentations
  • query view → \(q = f_q(\tilde{\mathbf{x}}_q)\)
  • key view → \(k_+ = f_k(\tilde{\mathbf{x}}_k)\)
• Maintain a queue of preview keys \(\{ k_i\}_{i=1}^K\) as negatives, where K can be large and is decoupled from batch size
  • FIFO queue so oldest (most likely stale) keys are removed when new keys added

Algorithm

InfoNCE over a dictionary

Cosine similarity via dot product (usually with normalized vectors):

Momentum encoder (EMA) update

m close to 1 so \(\theta_k\) doesn't vary too crazily

• if you unroll it, it follows exponential structure \(\theta_k(t)= (1 - m)\sum_{i=0}^{t-1} m^{i} \theta_q^{(t - i)}+ m^{t} \theta_k^{(0)}.\)

Contrastive learning optimized with

• many negatives
• consistent encoding of those negatives over time (instead of old embeddings)
  • uses EMA key encoder

Shuffling Batch Norm

• both encoders have Batch Norm
• Similar to ResNet paper, using BN prevents the model from learning good representations
  • cheats and find a low-loss solution as intra-batch communication causes leakage
    • If a query view and its positive key view land on the same GPU in the same step, they share the same BN statistics
• resolve by shuffling BN
  • query input encoder order unchanged
  • key encoder: randomly shuffle the minibatch before splitting across GPUs, run BN locally per GPU as usual, then unshuffle back afterward

Bootstrap your own latent: A new approach to self-supervised Learning

Core Modeling Recipe (BYOL)

• sample two augmented views of the sample image: v and v'
• Online network:
  • encoder f
  • projector g
  • predictor q
• Target network has encoder + projector but no predictor; updated by EMA from online params
• Stop gradient (sg) on target branch: target outputs are constants in the loss

One-direction loss: (predict target embedding)

Symmetric objective:

EMA update

Collapse? Why not constant representation?

• asymmetric as target doesn't have predictor
  • makes it not self-distillation
• EMA is a slow moving target

Emerging Properties in Self-Supervised Vision Transformers

Core Modeling recipe (DINO)

• make multiple augmented views of each image using multi-crop: a couple of global crops + several local crops
• run student on all views, run teacher on only global views
  • cheaper + stabilizes targets
• Student is trained with cross-entropy to predict the teacher's output distribution
  • teacher's branch is stopgrad
• Teacher parameters updated via EMA of student parameters (like MoCo and BYOL)

Stabilizers

• Sharpening (\(\tau\)): teacher uses a low temperature so targets are harsh
  • to enforce decisions

• Centering: subtract a running mean of teacher logits to prevent collapse to a constant distribution (applied on the teacher)
  • removes global bias direction as teacher can't encode info via a shared offset. Deviations from mean are what survives

Distillation Objective

image x, V(x) be augmented views. Student and teacher produce K-dimensional logit vectors

Teacher distribution:

Student distribution:

Loss: cross-entropy where teacher is the target

• outer iterate over teacher (global views)
• iterate over all views that are not same view as teacher to prevent trivial self-matching

Teacher EMA update:

Trained on VITs, learned representations/attention contain explicit object/region structure compared to supervised ViTs/convnets

• does well for k-NN and linear evaluation on ImageNet

This is because it forces global↔local consistency

Masked Autoencoders Are Scalable Vision Learners

Core modeling recipe (MAE)

• High masking ratio: mask a large fraction of patches (\(\sim\)75$\%$)
• Asymmetric encoder-decoder:
  • encoder sees only visible patches (no mask tokens)
  • decoder gets encoded visibles + mask tokens (plus positions) and reconstructs pixels
  • decoder can be much smaller; less than 10$\%$ computation per token vs encoder

If the encoder only processes \(\sim\)25$\%$ of tokens, and attention is quadratic, you save a lot of compute/memory.

Math

Let image be patchified into N patches (tokens). Randomly split indices into:

• visible set V
• masked set \(M\), with \(|M|\) large

Encoded visible tokens only:

where X is input feature and P is positional encoding

Decoder by inserting mask tokens (unshuffle to original spatial order) and decode

During decoding, the encoder outputs are concatenated with \(m+P_j\), where \(m_M\) is a learned parameter that is shared across all masked tokens. We then unshuffle to preserve original order.

MSE Loss (only on masked patches):

if you use reconstruction target is per-patch normalized pixels (0-255 → 0-1), representation quality improves

• so instead of matching the right RBG, its more about reconstructing local structure

if bottleneck is encoder FLOPs/memory, MAE is attractive as masking ratio can be high

• MAE supervision is reconstruction in pixel space, not distribution matching (like DINO)

iBOT: Image BERT Pre-Training with Online Tokenizer

• global semantics from \([\mathrm{CLS}]\) self-distillation (like DINO)
• local semantics via masked patch-token prediction (targets use teacher's online tokenizer)

Core modeling recipe (iBOT)

• take image x, sample two augmented views u, v. Apply blockwise masking to get masked views \(\hat{u}\), \(\hat{v}\).
  • teacher runs on 2 global crops
  • student runs on all crops (2 global + 10 local crops)
  • iBOT samples contiguous rectangular blocks
• Run student on masked view and teacher on unmasked view, teacher is an EMA of student with stop-grad
• Optimize two objectives
  • Cross-view \([\mathrm{CLS}]\) distillation
  • In-view masked patch-token distillation

Multi-crop setting

• 2 global crops (224 x 224)
• 10 local crops (96 x 96)

Math

• \([\mathrm{CLS}]\) self-distillation loss (like DINO's Cross Entropy)

where \(H(p,q) = -p^\intercal \log q\).

We take \([\mathrm{CLS}]\) embedding from view and pass through iBOT head to get probability vector over head's K outputs

• Masked image modeling loss

similarly with v. Teacher param \(\theta'\) are EMA of student params \(\theta\).

• teacher outputs soft token distributions per patch; student learns to “fill in” masked patches by matching those distributions
• Note this only over masked positions

Why online tokenizer?

• classic “BERT for images” require a separate pretrained tokenizer
  • iBOT uses the teacher network + projection head
• we care about soft token distribution (not one-hot)

for \([\mathrm{CLS}]\) and all patch tokens

• this is to force MIM (masked image modeling) to learn the global semantics

where \(H(p,q) = -p^\intercal \log q\).

TTT

TTT-E2E

• do test-time learning on the model weights, not on a cache or memory tokens. Context gets compressed into parameter updates
• loss is next-token-prediction at test time, so modeling is (inner loop is gradient step of this loss):

• only a subset of model weights are being updated during online gradient descent during inference
• its end-to-end because we use the same token-level LM loss for both outer and inner training

TTT-KVB

• test-time training framed as key-value binding: learn a map from keys to values in an inner loop
• KV reconstruction / regression-style loss: \(\phi\) is adapted at test time

• This is not best understood as literal memorization, as a broad class of TTT-KVB models can be rewritten as learned linear attention

Test-Time Training with KV Binding Is Secretly Linear Attention

GradMem

• Instead of changing full weights, just do test-time optimization on a small set of prefix memory tokens while freezing base model
• Loss: a self-supervised context reconstruction loss over memory tokens. We optimize \(M\) with a few GD steps

• Only memory-token parameters \(M\) is updated, not the backbone weights
• somewhat close to latent fitting / compressed memory writing rather than online weight adaptation
  • memory capacity bottlenecked by number of learned memory tokens

FALCON

• treat recurrent fast-weight state \(S_t\) as an online learner but using aligned pair \((k_{t-1}, v_t)\) ← predicting the future
• Regression objective:

• gradient step:

• NLMS normalization

• variants
  • FALCON-2: single-step online ridge regression with per-channel adaptive step sizes
  • FALCON-3: sliding window mini-batch regression
  • FALCON-3A: inner-product version with decoupled plasticity and forgetting
• FALCON is more close to a derived recurrent online optimizer

Optimizers

Adam

Adam forms exponential moving averages of the gradient and squared gradient:

With bias correction,

Adam then updates

where the division and square root are coordinatewise. So Adam is momentum with coordinatewise scale normalization as each parameter's step is normalized by its own recent gradient magnitude.

AdamW

AdamW uses the same adaptive Adam update, but decouples weight decay from the gradient update. In Adam with ordinary \(L_2\) regularization, the penalty is added into the gradient:

Then Adam normalizes this full gradient coordinatewise:

This means the decay term is also divided by Adam's adaptive denominator, so the effective shrinkage of \(W_{ij}\) depends on \(\sqrt{\hat v_{ij}}\). Parameters with large historical gradients get less decay, while parameters with small historical gradients get more decay.

AdamW fixes this by applying weight decay directly to the weights:

Equivalently,

So AdamW is Adam plus direct weight decay, where the decay term uniformly pulls weights toward zero instead of being distorted by Adam's coordinatewise normalization.

Muon

Muon forms a momentum matrix for each weight matrix:

momentum

Instead of applying \(M_t\) directly, Muon orthogonalizes its update direction. Taking

Muon replaces \(M_t\) by

Equivalently, it keeps the singular vectors of \(M_t\) but sets all singular values of the update direction to \(1\):

The update is then

Thus Muon is matrix-level momentum with spectral normalization of the update: it preserves the left/right singular directions of the gradient while removing their singular-value magnitudes.

Implementation

So for 2D parameters of neural network hidden layers, we define it as

Where 'NewtonSchulz5' is the Newton-Schulz matrix itteration:

def newtonschulz5(G, steps=5, eps=1e-7):
    a, b, c = (3.4445, -4.7750, 2.0315)
    X = G.bfloat16()
    X /= (X.norm() + eps)
    if G.size(0) > G.size(1):
        X = X.T
    for _ in range(steps):
        A = X @ X.T
        B = b * A + c * A @ A
        X = a * X + B @ X
    if G.size(0) > G.size(1):
        X = X.T
    return X

NS iteration orthogonalizes the update

Let $G = USV^\intercal$ be the SVD of the update matrix produced by SGD-momentum. One step of NS iteration w/ $(a,b,c)$ yields

So $N$ steps of NS with coefficients $(a,b,c)$ gives output $U \varphi^N(S) V^\intercal$ where $\varphi(x) = ax + bx^3 + cx^5$.

Using an ad-hoc gradient based approach, we end up with coefficients $(3.4445, -4.7750, 2.0315)$, which is the Keller Jordan initialization.

Muon vs Adam

At scale, Muon-trained transformers use QK-norm (or a clip like MuonClip) as it is more prone to get logit blowup (vs Adam)

• logits between tokens: $l_{ij} = \frac{q_i \cdot k_j}{\sqrt{d}}$, with $q_i = W_Q x_i$, $k_j = W_K x_j$.
• Bounding by spectral norms: \[l_{i,j} \lesssim \frac{\sigma_{\text{max}}(W_Q)\sigma_{\text{max}}(W_K)\Vert x_i \Vert \Vert x_j \Vert}{\sqrt{d}} = A \sigma_Q \sigma_K \sim \sigma_Q \sigma_K \sqrt{d}\]

  with $A = \frac{\Vert x_i \Vert \Vert x_j \Vert}{\sqrt{d}}$ or the Activation norm factor.

  NOTE The last approximation assumes $\Vert x_i \Vert \Vert x_j \Vert = d$ which occurs only with LayerNorm/RMSNorm (making per-coordinate RMS roughly 1). Usually holds for pre-norm attention.
• when softmax saturates with logit gap $\Delta l$ exceeds a few: attention goes one-hot, and gradient $\partial p / \partial l \propto p(1-p) \rightarrow 0$. So danger is $\sigma_Q \sigma_K$ getting large.
• Adam's saturation time is roughly the square of Muon saturation time, due to lack of spectral alignment.
  time of saturation (IMPORTANT)

  • For Muon, \[W_{t+1} = W_t - \eta O_t, \qquad O_t = U_t V_t^\intercal\]

    aka all singular values of $O_t$ are pushed to $1$. As a result, the top singular direction receives a step of size $\eta$ each time.

    Since $\Vert O_t\Vert_2=1$, each Muon step can change the spectral norm by at most $\eta$:

    \[ \sigma_{\max}(W_T)\le \sigma_{\max}(W_0)+\eta T. \]
    NOTE Here $\Vert O_t\Vert_2$ means the Schatten-$\infty$ norm / spectral norm: \[ \Vert O_t\Vert_{S_\infty} = \sigma_{\max}(O_t). \]

    Because Muon sets all nonzero singular values of the update to $1$,

    \[ \Vert O_t\Vert_{S_\infty}=1. \]

    More generally, the Frobenius norm upper-bounds the Schatten-$\infty$ norm:

    \[ \Vert A\Vert_{S_\infty} = \max_i \sigma_i(A) \le \left(\sum_i \sigma_i(A)^2\right)^{1/2} = \Vert A\Vert_F. \]

    So controlling Frobenius norm is a looser sufficient way to control spectral/operator scale. Muon is stronger for the update direction because it directly normalizes the singular values of the matrix update, giving a spectral-norm-controlled step.

    If the updates have persistent positive projection onto the top singular direction, with average alignment $c_Q$, then heuristically

    \[ \sigma_Q(T)\approx \sigma_Q(0)+c_Q\eta T. \] \[\sigma_K(T) \approx \sigma_K(0) + c_K\eta T\]

    with $0 \le c_Q, c_K \le 1$.

    Hence,

    \[\vert l(T) \vert \sim \sigma_Q(T) \sigma_K(T) \sim A c_Q c_K \eta^2 T^2\]

    Saturation occurs when $A c_Q c_K \eta^2 T^2 \sim C$ (threshold) or

    \[T_{\text{sat}}^\text{Muon} \sim \frac{1}{\eta} \sqrt{\frac{C\sqrt{d}}{c_Qc_K \Vert x_i \Vert \Vert x_j \Vert}} \sim \frac{1}{\eta d^{1/4}}\]

  • For Adam, our update $W_{t+1} = W_t - \eta \frac{m_t}{\sqrt{v_t} + \epsilon}$.

    Without orthogonalization, the growth of $\sigma_Q, \sigma_K$ is more naturally modeled as random-walk-like (asymptotic to $\sqrt{T}$):

    \[\sigma_Q(T) \sim \sigma_Q(0) + a_Q \eta \sqrt{T}\] \[\sigma_K(T) \sim \sigma_K(0) + a_K \eta \sqrt{T}\]

    Then

    \[\vert l(T) \vert \lesssim A \sigma_Q(T) \sigma_K(T) \sim A a_Q a_K \eta^2 T\]

    So

    \[T_{\text{sat}}^\text{Adam} \sim \frac{C\sqrt{d}}{a_Qa_K\eta^2 \Vert x_i \Vert \Vert x_j \Vert} \sim \frac{1}{\eta^2 d^{1/2}}\]

• With QK-norm we replace $q_i$ and $k_j$ with \[\tilde{q_i} = \alpha_q \frac{q_i}{\Vert q_k \Vert}, \qquad \tilde{k_j} = \alpha_k \frac{k_j}{\Vert k_j \Vert}\]

  This forces logit to be bounded by

  \[\vert l_{i,j} \vert = \Bigg \vert \frac{\tilde{q_i}^\intercal \tilde{k_j}}{\sqrt{d}} \Bigg \vert \le \frac{\alpha_q \alpha_k}{\sqrt{d}}\]
  NOTE In practice, many implementations use per-head LayerNorm/RMSNorm on q and k instead of L2Norm

MuonClip

Kimi K2

Muon is more token-efficient than AdamW but Muon scaling causes attention logit explosion more often.

QK-Clip: Per head:

If it exceeds a threshold $\tau$, then define

from activations $X$ in current training batch $B$.

We then rescale the projection weights after optimizer update:

which scales product $QK^\intercal$ by $\gamma_h$. Kimi K2 uses MLA-style attention, and they only clip the head-specific components (not the concat RoPE part).

algorithm

Gram Newton Schulz

[TODO]

Shampoo

[TODO]

Efficiency

TurboQuant

https://research.google/blog/turboquant-redefining-ai-efficiency-with-extreme-compression/

• traditional vector quantization usually introduces own “memory overhead” as most methods require calculating and storing (in FP) quantization constants for each small block of data
  • 1 - 2 extra bits per number which kind of defeat the purpose of vector quantization
    • scale (how far apart), zero-point (when values not centered), norm/block magnitude

Mechanism

• High-quality compression (PolarQuant): randomly rotating data vectors to simplify data's geometry
  • a random rotation already makes the coordinates statistically nice enough in high dimension
    • follows a Beta distribution
• TurboQuant uses a small, residual amount of compression power (1 bit) to apply QJL algorithm: acts as mathematical error-checker to eliminate bias

QJL

• Johnson-Lindenstrauss Transform to transform complex, high-dimensional data while preserving essential distances/relationships between data points
• reduces each vector number into a single sign bit (+/- 1)
• essentially just \(\text{sign}(Sx)\) where \(S\) is a random Johnson--Lindenstrauss-style random transform
• dequantization is \(\sqrt{\pi/2}\frac{1}{d}S^\intercal z\)
• important because plain quantization can preserve vector reasonably but still distort dot products in a biased way (important for attention)
  • 1-bit residual correction to remove bias as it makes inner-product estimate much more accurate

Polar Quant

• first apply random rotation so transforemd vector behaves close to Gaussian
• convert vector into polar coordinates using Cartesian coordinate system
  • (x,y) to \((r, \theta)\)
  • radius signifies how strong core data is, angle indicates data's direction
  • higher-dimensional vectors are represented by one radius + hierarchy of angles
• with this, model no longer needs to perform expensive data normalization step
  • boundaries are already known
• angles are quantized
  • level-1 angles are uniform on \([0, 2\pi)\), higher level angles lie in \([0, \pi/2)\) (concentrated around \(\pi/4\))
  • can quantize each angle independently with optimized codebooks
    • finite set of representative angle values per set and rounds each angle at each level to its nearest codeword in that set

Scaling Laws

[TODO: Chinchilla, compute-optimal training, overtraining, data quality, repeated data, scaling exponents]

muP/muTransfer

(Tuning Large Neural Networks via Zero-Shot Hyperparameter Transfer)

$\mu$P changes parameterization so that effect of each training update on activations stays $O(1)$ as width grows. Hence, we can

which is $\mu$Transfer.

$\mu$P instead scales updates so that $\Delta W_{ij} = O(1/n)$ so that

muP Principle

$\mu$P tries to make the learned function width-stable:

for models of different widths $n$, $N$ will assumed same hyperparameters.

For hidden weights $W \in \mathbb{R}^{n \times n}$, we want

and

so that

With SGD:

and $\mu$P chooses layerwise LR so that induced change is $O(1)$.

Scaling Laws, Carefully

Scaling Laws, Carefully

ML Loss Predictability

Observation is training loss $L$ decreases predictably as we scale up model size $N$, dataset size $D$, and compute $C$, following a power-law curve. This appears as a straight line on a log-log plot.

Notation: For reference,

• $N$: model size (in parameter count)
• $D$: training dataset size (in token count)
• $C$: training compute in FLOPs (Kaplan et al approximates as $C \approx 6ND$)
• $E$: irreducible loss (entropy of true data distribution)
• $L$, $\hat{L}(\cdot)$: test loss or its prediction function. Also refered to training loss since strongly correlated
• $\epsilon$: generalization error.

Usually

Early Days

Amari et al 1992 derived four types of learning curves using Bayesian approach and annealed approximation:

• Deterministic learning algorithm, noiseless data, one unique solution: $\epsilon \sim c \cdot D^{-1}$, for constant $c$
• Deterministic learning algorithm, noiseless data, multiple equivalent solutions: $\epsilon \sim c \cdot D^{-2}$. Model only learns the optimal manifold of parameters, not unique solution, so with more data, it learns faster
• Deterministic learning algorithm, noisy data: $\epsilon \sim c \cdot D^{-1/2}$ as with noise it is harder to learn
• Stochastic learning algorithm, noisy data: $\epsilon \sim c \cdot D^{-1} + E$, where $E$ is the irreducible loss a stochastic learner can not reduce further
  • ex: model runs out of capacity on large data

Hestness et al. 2017 analyzed for a given dataset size, identified the best-fit model size then plotted loss against dataset size. Across different domains in deep learning,

• generalization error scales as a power law across factors (e.g. data size)
• model improvements shift the error curve but doesn't change the power law exponent
• architecture changes the offset ($E$) of the power-law fit but does not change the exponent.
  IMPORTANT: Slope is a factor of the problem domain (not architecture)
• the number of model parameters $N$ needed to fit a dataset of size $D$ also scales as a power law.

3 stages

We can break the learning curve into 3 stages:

Rosenfeld et al. 2020 models error as a joint function of model size $N$ and dataset size $D$ across architectures (ResNet, Transformer, WRN, ...) and optimizers.

Holding one axis fixed, they found:

providing a joint form:

where $E$ is not dependent on $N,D$.

This is important as they can just build a prediction model using $\theta = \Big< A,B,E,\alpha,\beta\Big>$ to predict loss for $(D,N)$ above certain threshols using just smaller $(D,N)$.

visual

Scaling Laws in Data-Infinite Region

Kaplan et al 2020 found that cross-entropy loss $L$ scales as a power law relative to $N$ (model size, excluding embeding layers), dataset size $D$, and training compute across many orders of magnitudes $C$.

Key Findings:

• $L$ scales as a power law with $N$, $D$, and $C$ individually. Optimally, use them jointly.
• Training curves follow predictive power laws where parameters are roughly independent of model size.
• Larger models are more sample-efficient, meaning that they reach a given loss with fewer optimization steps and fewer data points than small models
• Architectural details (e.g. width, aspect ratio) matter less than sheer scale
• Train loss and test loss are positively correlated (this doesn't imply it transfers to post-training evaluation)
• Given a fixed compute budget, it is more efficient to train a very large model and stop before convergence than to train a smaller model all the way to convergence.

The joint dependence on $N$ and $D$ is

Approximation of FLOPs based on $D$ and $N$ (per multiply-add is $\sim 2$ FLOPs):

Scaling Laws in Data-Limited Region

Trickiness of Fitting Scaling Laws in Reality

Data

[TODO: deduplication, filtering, data mixtures, synthetic data, code/math data, contamination, curriculum]

Objectives

[TODO: next-token prediction, masked modeling, denoising, contrastive losses, reconstruction, preference losses, RL losses]

Pre-Training

[TODO: purpose of warmup rate, weight decay]

Post-Training

[TODO: SFT, RLHF, RLAIF, DPO/IPO/KTO, GRPO, PRMs, verifiers, tool-use RL, agent RL]

Evaluation

[TODO: perplexity, downstream benchmarks, long-context evals, code/math evals, agent benchmarks, contamination, reward hacking]

Tokenization

[TODO: BPE, SentencePiece, byte-level tokenization, multilingual fragmentation, image/audio/video tokens, latent tokens]

MoE

[TODO: Switch Transformer, GShard, Mixtral, DeepSeek MoE, routed/shared experts, load balancing, expert parallelism]

Memory

[TODO: KV cache, retrieval, recurrent memory, compressive memory, memory tokens, external memory, context extension]

Multimodal

[TODO: CLIP, Flamingo, BLIP, LLaVA, VLM training, video models, robotics policies, world models]

Reasoning and Agents

[TODO: chain-of-thought, self-consistency, search, verifiers, ReAct, tool use, planning, multi-agent systems]

Safety

[TODO: pretraining safety, refusal tuning, constitutional AI, scalable oversight, deception, eval awareness, jailbreaks, control]

Regularization and Stabilization

[TODO: RMSNorm, LayerNorm, QK-norm, residual scaling, clipping, weight decay, z-loss, logit clipping, warmup schedules]

Distillation

[TODO: logit distillation, feature distillation, sequence-level distillation, teacher-student setup, reasoning distillation]

Teacher-Student Compatibility

(Strong Teacher Not Needed? On Distillation in LLM Pretraining )

In LLM pretraining, distillation quality $\neq$ f(teacher strength only)

For sequence $x = (x_1, \dots, x_T)$, the student defines $p_{\theta_S}(\cdot \vert x_{\lt t})$ and teacher defines $p_{\theta_T}(\cdot \vert x_{\lt t})$.

Standard pretraining uses ground-truth next token:

which is a one-hot target as all probability mass is placed on observed token $x_t$.

Distillation instead uses teacher's full soft distribution. So,

with KL over all vocab $v \in V$.

Mixed training objective is

where $\alpha$ is amount of teacher supervision.

• Best teacher is not necessarily strongest teacher
  • distillation works when teacher's distribution gives useful information that the student would not have learned as easily form one-hot data alone.
• Weak teacher can still help if $\alpha$ is small (useful hints)

Inference Stack

[TODO: vLLM, SGLang, TGI, TensorRT-LLM, continuous batching, paged attention, prefix caching, chunked prefill, speculative decoding, KV cache management, latency vs throughput]

Reuse Middle Layers

Training-Free Looped Transformers

Original model:

Choose middle window $[a,b]$, define

Then replace we replace model with

Two ways:

• Block-mode:

  \[g_\text{blk}^{(K)}(x) = (L_b \circ \cdots \circ L_a)^K(x)\]

• Layer-mode:

  \[g_\text{lyr}^{(K)}(x) = L_b^K \circ L_{b-1}^K \circ \cdots \circ L_a^K(x)\]

Naive Looping Fails

A pre-norm transformer layer:

\[L(x) = x + \text{Attn}(\text{LN}_1(x)) + \text{MLP}(\text{LN}_2(x+\text{Attn}(\text{LN}_1(x))))\]

For a loop window $g$ its residual field is defined as $F_g(x) := g(x)-x$. This gives $g(x) = x + F_g(x)$.

We can interpret this as a forward Euler step with step size $h=1$ on

\[\dot{x} = F_g(x)\]

euler step

For ODE $y'(t)=f(t,y(t))$ with $y(t_0)=y_0$, \[y_{n+1} = y_n + hf(t_n, y_n)\] \[t_{n+1} = t_n + h\]

In naive looping: $x_{k+1} = g(x_k) = x_k +F_g(x_k)$ which when repeated $K$ times, approximates hidden state at $t=K$ not $t=1$.

So Naive Looping essentially increases effective depth by $K$ times, but the loop should still act as $1$ layer.

Hence, we take smaller steps for the same endpoint:

where the residual field is defined as $F_g(x) := g(x)-x$ for window $g$.

Algorithm.

Distributed

PCIe: the standard CPU/GPU expansion bus. GPUs, NICs, SSDs, etc. connect through PCIe lanes. It is general-purpose and widely available, but GPU-to-GPU communication often has to go through the PCIe switch/root complex, so bandwidth is lower and latency is higher than NVLink.

NVLink: NVIDIA's high-bandwidth GPU-to-GPU interconnect. It is designed specifically for fast accelerator communication, so GPUs can exchange tensors much faster than through ordinary PCIe. In multi-GPU training, NVLink helps NCCL all-reduce gradients/activations more efficiently inside a node.

[TODO: DDP, FSDP, ZeRO, tensor parallelism, pipeline parallelism, sequence parallelism, expert parallelism, HSDP, multi-cluster]

MultiKueue

C10d

gloo vs mpi vs nccl vs xccl

Branch-Train-Merge

Branch-Train-Merge trains many independent expert langauge models instead of one fully synchronized dense model.

Setup

Let the training corpus be split into domains:

where BTM trains an expert model for each domain:

where each $\theta_i$ is a complete language model specialized to the domain $D_i$.

The LM objective for expert $i$ is just standard next-token likelihood:

Algorithm

Step 0: Seed

Diffusion and Flow Models

[TODO: DDPM, score matching, flow matching, rectified flow, consistency models, diffusion transformers, video diffusion]

Async RL

[TODO: actor-learner systems, rollout workers, policy lag, stale gradients, asynchronous sampling, verifier latency, environment throughput, online RL at scale]

Continual Learning

Online Meta-Learning

Meta-Learning Online Adaptation of Language Models

Static LMs go stale. A natural fix is:

per new document $x$ in an online stream.

However naive online LM fine-tuning barely stores useful factual information.

• token NLL isn't aligned with token usefulness aka \[\nabla_\theta \mathcal{L}_{LM} = \sum_{t} \nabla_\theta [-\log p_\theta(x_t \vert x_{\lt t})]\]

  contains a mix of both useful factual gradients and noisy gradients.

CaMeLS (Context-aware Meta-learned Loss Scaling): learn a weighting model $w_\phi(x)$ that assigns each token an importance weight.

• weighting model is a small autoregressive model to decide which tokens are worth updating on

Weighted LM loss:

with

Meta-Learning Objective: Triples $(x,q,y)$ where

• $x$: document
• $q$: question about document
• $y$: answer

Inner loop updates $\theta$:

Outer loop updates $\phi$:

with loss being a sum over all positions $i \in \{1, \dots, T\}$:

• asks whether if the update avoided unnecessarily changing the model's general behavior

Findings:

The learned distribution is also sparse/bimodal where most tokens get low weights and a small set gets high weight (usually proper nouns/numbers).