22580: от GPT-2 до Kimi K3, объяснение

Полный перевод X Article о пути архитектур LLM от GPT-2 к Kimi K3: KV cache, linear attention, DeltaNet, KDA, MoE, MLA и AttnRes.

22580: от GPT-2 до Kimi K3, объяснение

Оригинал опубликован ali в X.

Двадцать две тысячи пятьсот восемьдесят. Именно столько моделей GPT-2 образца 2019 года помещается внутри Kimi K3 образца 2026 года. За семь лет масштаб вырос в 22 580 раз. Но дело только в масштабе?

В этом worklog я пройду путь до этой точки и покажу, насколько много или, наоборот, насколько мало на самом деле изменилось с тех пор. Мы проследим основные архитектурные изменения, которые привели к Kimi K3.

GPT-2

GPT-2 — decoder-only архитектура:

tok_emb = self.transformer.wte(idx) # token embeddings of shape (b, t, n_embd)
pos_emb = self.transformer.wpe(pos) # position embeddings of shape (t, n_embd)
x = self.transformer.drop(tok_emb + pos_emb)
for block in self.transformer.h:
    x = block(x)
x = self.transformer.ln_f(x)
logits = self.lm_head(x)
return logits

На вход подаются токенные и позиционные эмбеддинги:

Если приблизить каждый transformer block, он выглядит так:

class Block(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.ln_1 = LayerNorm(config.n_embd, bias=config.bias)
        self.attn = CausalSelfAttention(config)
        self.ln_2 = LayerNorm(config.n_embd, bias=config.bias)
        self.mlp = MLP(config)

    def forward(self, x):
        x = x + self.attn(self.ln_1(x))
        x = x + self.mlp(self.ln_2(x))
        return x

Процесс attention:

        B, T, C = x.size() # batch size, sequence length, embedding dimensionality (n_embd)

        # calculate query, key, values for all heads in batch and move head forward to be the batch dim
        q, k, v  = self.c_attn(x).split(self.n_embd, dim=2)
        k = k.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        q = q.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)
        v = v.view(B, T, self.n_head, C // self.n_head).transpose(1, 2) # (B, nh, T, hs)

        # manual implementation of attention
        att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
        att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
        att = F.softmax(att, dim=-1)
        att = self.attn_dropout(att)
        y = att @ v # (B, nh, T, T) x (B, nh, T, hs) -> (B, nh, T, hs)
        y = y.transpose(1, 2).contiguous().view(B, T, C) # re-assemble all head outputs side by side

        # output projection
        y = self.resid_dropout(self.c_proj(y))
        return y

Когда получена финальная матрица hidden states, language-model head отображает ее в logits по словарю. При autoregressive decoding для выбора следующего токена нужны только logits последней позиции.

Это неэффективность decoder-only генерации: модель вычисляет представления для каждой входной позиции, но на каждом шаге декодирования использует только logits последней позиции. Без кэша значительная часть работы повторялась бы для следующего токена.

KV cache следует из простого наблюдения: после добавления сгенерированного токена к входу модель иначе заново вычисляла бы projections для всех предыдущих токенов. Сохранение их key и value vectors устраняет эту лишнюю работу.

Это хранилище и есть KV cache. Оно удерживает vectors для предыдущих N-1 токенов и может стать настолько большим, что упирается в пропускную способность памяти.

В целом, при примерно 50 тыс. возможных токенов, 12 блоках, 12 heads и embedding dimension 768 базовая модель содержит около 124 млн параметров.

vocab_size: int = 50304 # GPT-2 vocab_size of 50257, padded up to nearest multiple of 64 for efficiency
n_layer: int = 12
n_head: int = 12
n_embd: int = 768

При 2,8 трлн параметров одна модель Kimi K3 содержит примерно столько же параметров, сколько 22 580 моделей GPT-2.

Linear Attention

Softmax attention применяет нелинейность после произведения q·k, связывая каждый query с каждым key. Linear attention вместо этого применяет feature map, например ELU+1, к q и k по отдельности. Это позволяет переассоциировать произведение, так что растущее множество K и V vectors можно свернуть в фиксированное состояние D×D.

Формулировка O(N²) в статье сначала сбила меня с толку. Неверно, что «стоимость одного time-step у трансформеров растет с квадратом текущей длины последовательности». Именно это исправляет FlashAttention... а потом я увидел, что он вышел в 2020 году.

В то время при обучении часто материализовали полную матрицу attention N×N, FlashAttention еще не существовал, а reference autoregressive implementations нередко пересчитывали историю токенов без KV cache.

def forward(self, x, mask=None, past_kv=None):
  # x is b,t,d
  b,t,d=x.shape
  d_head=d//self.num_heads
  h=self.num_heads
  qkv=self.qkv_proj(x)

  q=qkv[:, :, :d].view(b,t,h,d_head).transpose(1,2)
  k=qkv[:, :, d:2*d].view(b,t,h,d_head).transpose(1,2)
  v=qkv[:, :, 2*d:].view(b,t,h,d_head).transpose(1,2)

  # at prefill, q,k,v have shapes b,h,t,d
  # at decode, shape is b, h, 1, d
  # so i cat at the t dimension, dim(2)

  if past_kv is not None:
    k_past=past_kv[0]
    v_past=past_kv[1]
    k=torch.cat((k_past, k), dim=2)
    v=torch.cat((v_past, v), dim=2)

  scores=(q@k.transpose(-1,-2))/math.sqrt(d_head)
  if past_kv is None: #we're in prefill and need to mask
    causal_mask=torch.ones(t,t,dtype=bool, device=q.device)
    causal_mask=torch.triu(causal_mask, diagonal=1)
    scores=scores.masked_fill(causal_mask, float('-inf'))

  if mask is not None:
    scores=scores.masked_fill(~mask, float('-inf'))

  #get attn (bhtt x bhtd)
  attn=scores.softmax(-1)#bhtt
  o=attn@v #bhtd
  o=o.transpose(1,2).contiguous().view(b,t,d)  #b,t,d

  # use x to get qkv
  o_proj=self.o_proj(o)
  past_kv=(k, v)
  return o_proj, past_kv

Тот же процесс проще увидеть визуально. Каждый шаг decode выполняет два ND-чтения и две 1D-записи в HBM, а KV cache растет линейно, O(N), вместе с длиной последовательности.

Обратите внимание на избыточные чтения и записи. В этой работе они заменяются на:

def forward(self, x, mask=None, cache=None):
  # x is b,t,d
  b,t,d=x.shape
  d_head=d//self.num_heads
  h=self.num_heads
  qkv=self.qkv_proj(x)

  q=qkv[:, :, :d].view(b,t,h,d_head).transpose(1,2)
  k=qkv[:, :, d:2*d].view(b,t,h,d_head).transpose(1,2)
  v=qkv[:, :, 2*d:].view(b,t,h,d_head).transpose(1,2)
  
  k=F.elu(k)+1 
  k=k.transpose(-1,-2) 
  q=F.elu(q)+1
 
  S,z=cache if cache is not None else (0.0, 0.0)
  S=S+k@v
  z=z+k
      
 o=q@S #bhtd
 denom=q@z
 o_scaled=o/denom
 o_scaled=o_scaled.transpose(1,2).contiguous().view(b,t,d)
 o_proj=self.o_proj(o_scaled)
 cache=(S,z)
 
 return o_proj, cache

Здесь есть компромисс.

Мы заменяем exponent, который использует softmax, на ELU+1, примененный отдельно к q и k до их взаимодействия. Оба подхода нормализуют получившиеся scores, но feature map в linear attention — менее выразительная аппроксимация softmax kernel. Такая аппроксимация может снижать fidelity, хотя практическая потеря точности зависит от архитектуры и workload.

Заметьте, что мы все еще делим на сумму qk, хотя на диаграмме это опущено для простоты. На высоком уровне attention состоит из трех шагов:

  1. Сделать qk scores неотрицательными. Linear attention использует ELU+1, softmax использует exponentiation.
  2. Разделить на сумму.
  3. Вычислить weighted average значений.

Это сохраняет базовый контракт attention, но использует менее выразительный feature map, чтобы сделать QK scores неотрицательными.

DeltaNet, или Fast Weight Programmers

Конечный cache должен перезаписывать информацию или объединять ее с уже сохраненной. State от токена i-1 не получает собственного слота; он добавляется в ту же матрицу D×D. Поэтому новые queries уже не могут достать идеально изолированное представление каждого более раннего токена.

Именно это добавление и дает выигрыш в эффективности. Аддитивное обновление cache вместо concatenation не дает ему расти как O(N), но та же операция вызывает interference информации. DeltaNet решает эту потерю recoverability.

Schlag в статье Fast Weight Programmers формулирует это так: «когда длина последовательности превышает емкость хранилища, модель может оказаться в режиме overcapacity. Чтобы корректно работать в таком режиме, модель должна научиться динамически взаимодействовать с содержимым памяти и избирательно решать, какие key-value associations сохранить, а какие удалить. Чисто аддитивная инструкция может быть неподходящей для этой цели... бесконечное добавление новых associations в память конечного размера, как в Eq. 17, неизбежно достигнет предела».

Режим, делающий linear attention привлекательным, где N намного больше D, одновременно вскрывает его главное ограничение. Когда state превышает эффективную емкость, associations начинают мешать друг другу, потому что update аддитивный и из cache ничего не уходит.

def forward(self, x, mask=None, cache=None):
  # x is b,t,d
  b,t,d=x.shape
  d_head=d//self.num_heads
  h=self.num_heads
  qkv=self.qkv_proj(x)

  q=qkv[:, :, :d].view(b,t,h,d_head).transpose(1,2)
  k=qkv[:, :, d:2*d].view(b,t,h,d_head).transpose(1,2)
  v=qkv[:, :, 2*d:].view(b,t,h,d_head).transpose(1,2)

  q = F.normalize(F.silu(q), dim=-1)     
  k = F.normalize(F.silu(k), dim=-1)     
  beta = torch.sigmoid(self.w_beta(x)).view(b, 1, t, 1)   
  # new: per-token write strength

  S = cache if cache is not None else 0.0  

  v_old = k @ S # read the board at this key
  u = beta * (v - v_old) # the delta: only what's actually new
  S = S + k.transpose(-1, -2) @ u # same outer-product write as before

  o = q @ S # read, no denominator
  o = o.transpose(1, 2).contiguous().view(b, t, d)
  return self.o_proj(o), S

Визуальный пример упрощает понимание.

Возьмем одну association, записанную как S = k.T @ v. Если читать ее тем же key, получаем k @ (k.T @ v), то есть (k @ k.T) v — squared norm key, умноженный на v. Поэтому read возвращает v, масштабированное squared norm key. Если нормализовать k до unit length или просто разделить результат на norm, получим v обратно точно.

Q — тоже learned pointer. Wq и Wk читают один и тот же residual stream, а query для факта указывает на key direction, в которую этот факт был записан. Update сначала спрашивает, какая информация сейчас извлекается из cache по текущему key. Он вычитает эту старую информацию из value, который мы хотим сохранить, умножает key на разницу и добавляет результат обратно. Старая информация удаляется, новая записывается на ее место.

DeltaNet, или Parallelizing Linear Transformers with Delta Rule

Это самый сложный раздел поста. Мне потребовалось около семи часов, чтобы выстроить рабочее понимание, поэтому объяснение пойдет от implementation. Коротко: DeltaNet реализует first-order linear recurrence с generalized Householder transition matrices, что позволяет выполнять chunk-wise parallel forward passes для hardware-efficient linear-time training. Он делит inputs и outputs на несколько chunks размера C и вычисляет outputs для каждого chunk на основе final state предыдущего chunk и query/key/value blocks текущего chunk.

Практическая проблема — prefill. Прямая реализация Delta rule по последовательности из T токенов выглядела бы так:

S = torch.zeros(b, h, dh, dh) if cache is None else cache
outs = []
for i in range(t):
    k_i = k[:, :, i:i+1]  
    v_i = v[:, :, i:i+1]
    b_i = beta[:, :, i:i+1]
    v_old = k_i @ S                  
    u_i  = b_i * (v_i - v_old)
    S = S + k_i.transpose(-1, -2) @ u_i # write
    outs.append(q[:, :, i:i+1] @ S)     
o = torch.cat(outs, dim=2)

В отличие от стандартного attention, такая формулировка требует correction на каждом key vector, поэтому путь к parallel matrix multiplication не сразу очевиден. Даже без Delta rule прямой prefill для linear attention остается последовательным:

S = torch.zeros(b, h, dh, dh) if cache is None else cache
outs = []
for i in range(t):
    q = q[:, :, i:i+1]  
    k = k[:, :, i:i+1]  
    v = v[:, :, i:i+1]

    S=S_old+k@v
	  o=q@S #bhtd
	  o=self.norm(o)
    o=o.transpose(1, 2).contiguous().view(b, t, d)

    out=self.o_proj(o)
    cache=S
    outs.append(out)

o = torch.cat(outs, dim=2)

Chunked formulation дает более эффективный подход. Механику проще понять на примере:

Если C=N, мы получаем стандартное O(N²) attention; если C=1 — обычное linear attention. Промежуточные значения интерполируют между ними: мы обмениваем дополнительную работу внутри chunk на лучшую загрузку hardware. На практике C часто равен 64 или 128, потому что tensor-core instructions эффективно работают на такой гранулярности; UMMA — один из примеров.

Промежуточные tiles сворачиваются в S как часть state update:

S = torch.zeros(b, h, dh, dh) if cache is None else cache
outs = []
for i in range(t//C):
    q_c = q[:, :, i*C:(i+1)*C]  
    k_c = k[:, :, i*C:(i+1)*C]  
    v_c = v[:, :, i*C:(i+1)*C]

	  o_prev=q_c@S #this is everything up to this block
	  
	  attn=(q_c@k_c.transpose(-1,-2)).tril() #masked attention 
	  o_curr=attn@v_c
		  
		o=o_prev+o_curr
    
    S_new=k_c.transpose(-1,-2)@v_c #recurrent attention 
    S=S+S_new
    outs.append(o)

o = torch.cat(outs, dim=2)

Внутри блока мы делаем q(kᵀv). Это сначала scores — обычный порядок attention с masking. Между блоками мы следуем (kᵀv)q, то есть делаем recurrent order: сначала state. Attention растет как O(N²), а этот подход — нет. Внутри блока я делаю настоящее attention: masked QKᵀ умножается на V. Между блоками я сворачиваю все в state и читаю его обратно одним matmul. Поэтому cost распадается на две части. Есть фиксированная часть, 2Ld², это state work, ей безразличен C. И есть растущая часть, 2LCd, это score matrices на диагонали. Full attention — это случай C=L, и тогда второй член становится 2L²d, квадратичным. Чем меньше C, тем меньше FLOPs.

C=1 — самый дешевый вариант по чистым FLOPs, но не обязательно самый быстрый по wall-clock. GPU может выполнить больше арифметики быстрее, если работа хорошо ложится на matrix-multiply hardware.

Следующий шаг — распространить тот же подход на DeltaNet.

Базовая проблема проста: chunking method, используемый для чисто additive attention, напрямую не применяется к delta updates:

v_old = k_i @ S                  
u_i  = b_i * (v_i - v_old)

Нам нужен каждый отдельный state по порядку, чтобы вычислить информацию, которую требуется вычесть. Без математической re-parameterization это нельзя распараллелить тем же способом. Поэтому авторы переписывают delta updates из:

u=v_new-v_old
S_t= S_(t-1)+K.T@u
o=q@S_T

Здесь последовательный loop вычисляет одну delta за итерацию. Reparameterized form выглядит так:

S_t = S_{t-1}(I − β_t k_t k_tᵀ)  +  β_t v_t k_tᵀ
o_t = S_t q_t

Эта формулировка позволяет chunked code вычислять все C deltas сразу:

def chunk_delta_rule_forward(Q, K, V, beta, C):
		# L: sequence length, d: head dimension
		L, d = Q.shape
		# chunking
		Q, K, V = map(lambda x: x.reshape(-1,C,d), [Q, K, V])
		beta = beta.reshape(-1, C)
		K_beta = K * beta.unsqueeze(-1)
		V_beta = V * beta.unsqueeze(-1)
		
		# compute eq. 10 with vectorized forward substitution for fast inverse
		T = -(K_beta @ K.t()).tril(-1)
		for i in range(1, C):
				T[i, :i] = T[i, :i] + (T[i, :, None] * T[:, :i]).sum(-2)
		
		T += torch.eye(C)
		W = T @ K_beta
		U = T @ V_beta

		# chunkwise parallel. Eq. 8-9
		S = torch.zeros(d, d)
		O = torch.empty_like(V)
		
		for i in range(L//C):
				q_i, k_i, w_i = Q[i], K[i], W[i]
				u_i = U[i] - w_i @ S # the corrections, all of one chunk
				o_inter = q_i @ S
				A_i = (q_i @ k_i.t()).tril() #qk.t
				o_intra = A_i @ u_i # attention @ v (with corrections, so u)
				S += k_i.t() @ u_i # update state with addition 
				O[i] = o_intra + o_inter #update output with flash + recurrent
		return O.reshape(L, d)

Это приводит нас к первой точке сравнения: MHA против DeltaNet Transformers.

Gated Delta Net

Теперь у нас есть метод точечно изменять cache. С каждым новым фактом, то есть каждым новым key vector, мы можем посмотреть, какая старая информация хранится в этой точке, и заменить ее новой информацией, к которой хотим обращаться.

Однако такой механизм может забыть только association, для которой есть конкретная replacement. Он не умеет эффективно очищать несколько associations при context switch или в целом ослаблять memory, чтобы освободить capacity.

Если бы мы работали с чисто additive linear attention, добавить способность забывать было бы просто. Нужен был бы только параметр, управляющий forgetful state:

S_old=cache
S_new=k@v
# cache=S_old+S_new
cache=alpha * S_old + S_new

Это вклад Mamba-2. Мы ослабляем предыдущий cache, затем добавляем новый cache в полную силу, не позволяя state расти без границ.

Равномерно ослаблять все key-value associations на каждом time step динамическим коэффициентом — рабочий подход, и именно это делает Mamba. Но он не учитывает разную важность разных key-value associations.

То есть если модели нужно забыть одну конкретную association, все associations забываются одинаково. Delta rule, наоборот, умеет обновить один факт, но не имеет способа заставить остальные факты затухать.

Поэтому Gated Delta rule объединяет gated update rule из Mamba с Delta rule. Он добавляет параметр alpha, который при значении один переключается на чистый Delta rule, а при значении ноль очищает память. Сложность — реализовать это тем же методом parallel chunks.

Implementation использует ту же DeltaNet reparameterization, что описана в предыдущем разделе. Математика почти идентична, но добавляется data-dependent scalar от нуля до единицы, управляющий decay предыдущего state. Это объединяет эффективное обучение key-value associations с adaptive memory management.

Соответствующие изменения в коде показаны ниже:

Член γʳ/γⁱ учитывает cumulative decay. Токен, записанный на time step x и прочитанный на x+t, был умножен на αₓαₓ₊₁αₓ₊₂…αₓ₊ₜ. Это multiplicative analogue вычисления prefix-sum.

Получившаяся архитектура выглядит так:

KDA / Kimi Linear

На этом этапе исследователи начали экспериментировать с hybrid models, которые объединяют несколько видов attention внутри одной архитектуры, например Gated DeltaNet с Mamba.

Kimi Linear привлекла внимание одним центральным заявлением: в controlled comparisons она превзошла full attention. Авторы представили ее как drop-in architectural replacement с лучшим качеством и до 6x более высоким decode throughput.

Kimi Linear улучшает Gated DeltaNet, вводя fine-grained gating. Вместо одного scalar decay модель учит отдельное значение decay для каждого channel.

KDA update rule остается похожим, но код теперь выглядит примерно так:

Здесь alpha.reshape(nb, C, d) отражает самый значимый вклад статьи: тонкий контроль memory decay по каналам.

Если поставить Kimi Linear рядом с DeltaNet Transformer, архитектура Kimi Linear вносит три крупных изменения:

  1. Использует hybrid system, в которой чередуются слои Multi-head Latent Attention (MLA).
  2. Заменяет MLP на слой Mixture-of-Experts (MoE).
  3. Добавляет capacity в DeltaNet через alpha projection.

Более поздние разделы подробнее разбирают MLA и MoE. Пока важен главный вывод: это не слепое масштабирование. Дополнительная capacity имеет конкретную математическую цель: per-channel scale дает модели более тонкий контроль decay памяти.

Scaling laws остаются важны, но capacity нужно добавлять в правильное место и в форме, которую система может использовать. Каждая архитектура в этой последовательности добавляет capacity, чтобы закрыть конкретное ограничение предыдущей системы.

Kimi K3

В итоге language backbone Kimi K3 похож на модель Kimi Linear выше. Он содержит 23 четырехслойных macrocycles. В каждом macrocycle три слоя используют Kimi Delta Attention, а четвертый — Multi-head Latent Attention. Первый слой использует dense feed-forward network; все остальные слои используют latent Mixture-of-Experts.

На первый взгляд изменения относительно Kimi Linear выглядят умеренными:

  • Существенное увеличение масштаба
  • Blockwise AttnRes каждые 12 слоев
  • MLA query LoRA и output gating
  • Latent-space MoE
  • SiTU activations
  • Gated MLA

KDA дает constant-state recurrent memory, а периодические MLA layers сохраняют full softmax retrieval по контексту. Следующая упрощенная визуализация служит полезной опорой для изменений ниже.

Начнем с более прямых изменений: Gated MLA, latent-space MoE и SiTU activations.

Gated MLA определяет, какая часть каждой retrieved feature проходит из MLA в residual stream. Он делает это через element-wise multiplication с gate, спроецированным из input.

В обычном MoE learned router использует dot-product similarity, чтобы отправить каждый token в subset expert networks. В Kimi K3 всего 898 experts. Два shared experts обрабатывают каждый token; из оставшихся 896 router выбирает 16 для каждого token.

Kimi K3 также меняет activation у experts. Вместо применения SiLU к up projection, element-wise multiplication на gate и затем down projection он использует SiTU:

d = x.shape[-1] // 2
gate = x[..., :d].to(torch.float32)
up = x[..., d:].to(torch.float32)
situ_a = self.beta * torch.tanh(gate / self.beta) * torch.sigmoid(gate)
if self.linear_beta is not None:
    up = self.linear_beta * torch.tanh(up / self.linear_beta)
return (situ_a * up).to(x.dtype)

Модель также down-projects inputs в shared experts и up-projects их финальную сумму:

Это иллюстрирует повторяющуюся проблему inference. Без fused kernel новая activation почти в 3 раза медленнее оригинального пути. Компенсирующая оптимизация в том, что experts работают в compressed latent space, что делает их forward pass намного быстрее и почти вдвое снижает FLOPs.

Оставшиеся изменения — MLA query LoRA, output gating и blockwise Attention Residuals каждые 12 слоев. AttnRes добавляет примерно 2% latency при inference, но дает два важных преимущества:

  • Selective retrieval более ранних representations, что смягчает residual dilution и рост hidden state
  • Compute advantage 1.25x

AttnRes и MLA решают одно и то же базовое ограничение с разных сторон. KDA layers работают с state постоянного размера и неизбежно должны отбрасывать информацию. MLA retrieves из token context, а AttnRes retrieves из более ранних depth-wise representations.

AttnRes

Спасибо @chloey3k за помощь с этим разделом. В каждом forward pass input проходит через stack слоев. Здесь каждый слой состоит из attention block (KDA или MLA) и MLP или MoE block. Обычно input каждого слоя — это сумма original embedding и outputs всех предыдущих слоев, причем все они имеют одинаковый вес.

Здесь h_i — input слоя i, h_1 — embedding текущего token, то есть последнего token в последовательности на данный момент, а f_i(h_i) — output слоя i: attention или MLP block.

Проблема в отсутствии selective access. Разные типы слоев получают один и тот же aggregated state, хотя им могли бы быть полезны разные weightings. Поскольку recurrence чисто аддитивная, более поздние слои также должны учить все более крупные outputs, чтобы влиять на accumulated residual, и это может дестабилизировать training. Вместо того чтобы считать все слои равными, AttnRes умножает каждый term этой суммы на специализированный weight, позволяя модели отдавать больше важности тем слоям, которые полезнее в текущем контексте.

Каждый weight alpha_i вычисляется из query-key dot product. Query learned для каждого слоя, а keys и values приходят из более ранних residual-stream states. Scores нормализуются так, чтобы суммироваться в единицу, а затем используются для формирования weighted combination этих states.

Поэтому модели не нужно опираться только на непосредственного предшественника. AttnRes дает каждому слою selective access к outputs более ранних слоев, позволяя learned query доставать representations, наиболее полезные для текущего вычисления.

Псевдокод ниже применяет ту же идею на уровне block granularity. Block — это element-wise sum outputs attention и MLP, накопленных за 12 decoder layers и сохраненных как одно depth representation для дальнейшего AttnRes mixing.

Применять residual attention на каждом слое было бы слишком дорого для training и inference. Применение только на фиксированных block boundaries сохраняет большую часть пользы при меньшей цене. В Kimi K3 такая граница находится после каждых 12 decoder layers. За 23 четырехслойных macrocycles это дает восемь AttnRes blocks, что увеличивает скорость inference.

Возможно, самая важная часть функции block_attn_res выглядит так:

V = torch.stack(blocks + [partial_block]) # [N+1, B, T, D]
K = norm(V)
logits = torch.einsum('d, n b t d -> n b t', proj.weight.squeeze(), K)
h = torch.einsum('n b t, n b t d -> b t d', logits.softmax(0), V)
return h

На этом завершается путь от GPT-2 к Kimi K3.

Главное изменение — не один только масштаб. Каждый архитектурный шаг меняет то, что модель хранит, как она обновляет state или как извлекает информацию, которую fixed-size state не может сохранить.

Kimi K3 объединяет constant-state recurrent memory, периодический softmax retrieval, sparse expert capacity и selective depth-wise residual access. В результате система тратит дополнительную capacity там, где у нее есть конкретная функциональная роль.

По сути, associative memory фиксированной емкости, то есть фиксированных dimensions, нуждается в eviction policy: чисто additive linear operation со временем начинает добавлять interference при достижении capacity. Поэтому learned selection — gating, routing или decay — становится необходимой, а attention остается самым эффективным механизмом selective read.

Subscribe to Temperature 0.7 - AI блог об AI и роботах

Don’t miss out on the latest issues. Sign up now to get access to the library of members-only issues.
jamie@example.com
Subscribe