От учителя к студенту: путь к изотропным эмбеддингам в 29M модели. Часть 4/10
Важное предупреждение:
Весь код по моей методологии написан нейросетью DeepSeek, сам бы не осилил (я не программист).
За мной лишь отладка, исправление неточностей и ошибок.
4.1. Архитектура StudentTiny
Три компонента:
Input: input_ids [B, T], attention_mask [B, T]
↓
[Turbo encoder — 3 слоя, 312 dim]
↓ hidden [B, T, 312]
[Multi-head Attention Pooling — K=4]
↓ pooled [B, 312]
[Linear(312 → 768), bias=True]
↓ projected [B, 768]
|
┌── projected=False → [L2-normalize] → 312-dim (продакшен)
└── projected=True → [Linear(312→768)] → [L2-normalize] → 768-dim (дистилляция)Ключевая идея — два режима через один метод:
student(input_ids, attention_mask, projected=False) # [B, 312]
student(input_ids, attention_mask, projected=True) # [B, 768]
projected=False— то, что пойдёт в продакшен. Вектор 312-dim, L2-норм.projected=True— то, что сравнивается с учителем во время дистилляции. Вектор 768-dim, L2-норм.
Так не нужно плодить модели — одна, но с двумя выходами. Проекция останется в экспорте как "опциональный 768-dim режим".
Детали каждого компонента — ниже: pooling (4.2), проекция (4.3).
4.2. Attention Pooling с 4 головами
Зачем нужен pooling? В BERT-выходе last_hidden_state — это матрица [B, T, H]: по одному вектору на каждый токен. Нам же нужен один вектор на всё предложение. Классический способ — mean pooling (усреднить все токены). Мы делаем иначе: взвешенное усреднение с обучаемыми весами.
Почему 4 головы? Одна голова даёт один "взгляд" на текст — например, "акцентируй ключевые термины". Четыре головы дают четыре разных взгляда: одна может смотреть на термины, другая — на структуру, третья — на именованные сущности. Результаты усредняются.
Стоимость:
Параметры: 4 x 312 = 1,248 (пренебрежимо).
Скорость: 4 головы вместо 1. По нашим замерам — оверхэд <1%.
Код: +5 строк.
Формула:
scores[b, k, t] = <query[k], hidden[b, t]> / sqrt(H)
scores[b, k, t] = mask -inf (если t — padding, через masked_fill)
weights[b, k, t] = softmax(scores[b, k, :])
pooled[b, k, h] = sum_t weights[b, k, t] * hidden[b, t, h]
output[b, h] = mean_k pooled[b, k, h]
Размер выхода не растёт — 4 головы усредняются в один вектор 312-dim.
Sanity: маскирование padding обязательно. Иначе softmax будет "смотреть" на пустые токены.
4.3. Проекция 312→768
Зачем? Учитель даёт эмбеддинги 768-dim. Наш студент — 312-dim. Чтобы сравнить их напрямую, нужен линейный слой, отображающий 312-dim в 768-dim.
Формула:
projected = W @ pooled + b
W: [768, 312], b: [768]
Bias не обязателен: 768 параметров — пренебрежимо, а на старте может помочь быстрее сойтись. Не вредит — оставляем bias=True.
4.4. Скрипт сборки StudentTiny
Файл: scripts/4.1_build_student.py
"""StudentTiny: turbo + AttentionPooling (K голов) + проекция 312→768.
Собирает студента, сохраняет state_dict + токенизатор + meta.
"""
import json
import time
from pathlib import Path
import torch
import torch.nn as nn
import torch.nn.functional as F
import transformers
from transformers import AutoModel, AutoTokenizer
ROOT = Path(__file__).resolve().parent.parent
DATA = ROOT / "data"
META_DIR = DATA / "meta"
META_DIR.mkdir(parents=True, exist_ok=True)
ENCODER = "sergeyzh/rubert-tiny-turbo"
HIDDEN = 312
TEACHER_DIM = 768
NUM_HEADS = 4
OUT_MODEL = DATA / "student_init.pt"
OUT_TOKENIZER = DATA / "student_tokenizer"
META = META_DIR / "4.1_build_student_meta.json"
# Слои, которые будут разморожены в Phase B (верхние 2 из 3)
UNFREEZE_LAYERS = [1, 2]
SEED = 42
class AttentionPooling(nn.Module):
def __init__(self, hidden_size: int, num_heads: int = 4):
super().__init__()
self.query = nn.Parameter(torch.randn(num_heads, hidden_size) * 0.02)
self.scale = hidden_size ** -0.5
self.num_heads = num_heads
def forward(self, hidden, attention_mask):
# hidden: [B, T, H], mask: [B, T]
scores = torch.einsum("bth,kh->bkt", hidden, self.query) * self.scale
scores = scores.masked_fill(attention_mask[:, None, :] == 0, float("-inf"))
weights = torch.softmax(scores, dim=-1)
pooled = torch.einsum("bkt,bth->bkh", weights, hidden)
return pooled.mean(dim=1) # усреднение по головам
class StudentTiny(nn.Module):
def __init__(self, encoder, hidden: int = HIDDEN,
teacher_dim: int = TEACHER_DIM, num_heads: int = NUM_HEADS):
super().__init__()
self.encoder = encoder
self.pooling = AttentionPooling(hidden, num_heads)
self.projection = nn.Linear(hidden, teacher_dim) # bias=True
def forward(self, input_ids, attention_mask, projected: bool = False):
hidden = self.encoder(
input_ids=input_ids, attention_mask=attention_mask
).last_hidden_state
pooled = self.pooling(hidden, attention_mask)
if projected:
return F.normalize(self.projection(pooled), dim=-1)
return F.normalize(pooled, dim=-1)
def count_params_by_group(student):
return {
"encoder": sum(p.numel() for p in student.encoder.parameters()),
"pooling": sum(p.numel() for p in student.pooling.parameters()),
"projection": sum(p.numel() for p in student.projection.parameters()),
"total": sum(p.numel() for p in student.parameters()),
}
def check_unfreeze_layers(student, layer_indices):
"""Проверяет, что пути encoder.encoder.layer.{i}.* существуют, считает параметры."""
result = {}
for i in layer_indices:
prefix = f"encoder.encoder.layer.{i}."
params = [(n, p.numel()) for n, p in student.named_parameters()
if n.startswith(prefix)]
result[f"layer.{i}"] = {
"found": len(params) > 0,
"n_params": sum(p for _, p in params),
"n_tensors": len(params),
}
return result
def sanity_checks(encoder):
"""Проверки конфига энкодера: hidden_size, vocab, max_position_embeddings."""
cfg = encoder.config
return {
"hidden_size": {
"value": cfg.hidden_size,
"expected": HIDDEN,
"ok": cfg.hidden_size == HIDDEN,
},
"vocab_size": {
"value": cfg.vocab_size,
"expected": 83828,
"ok": cfg.vocab_size == 83828,
},
"max_position_embeddings": {
"value": cfg.max_position_embeddings,
"expected_min": 256,
"ok": cfg.max_position_embeddings >= 256,
},
"num_hidden_layers": {
"value": cfg.num_hidden_layers,
"ok": cfg.num_hidden_layers == 3,
},
}
def main():
t0 = time.time()
# Детерминированная инициализация query
torch.manual_seed(SEED)
print(f"Loading encoder: {ENCODER}", flush=True)
encoder = AutoModel.from_pretrained(ENCODER)
student = StudentTiny(encoder)
# 1. Параметры
params = count_params_by_group(student)
print(f"\nParams:")
print(f" encoder: {params['encoder']:>12,}")
print(f" pooling: {params['pooling']:>12,} (K={NUM_HEADS})")
print(f" projection: {params['projection']:>12,} (312→768, bias=True)")
print(f" TOTAL: {params['total']:>12,}")
# 2. Sanity checks
sanity = sanity_checks(encoder)
print(f"\nSanity checks:")
for name, info in sanity.items():
marker = "✓" if info["ok"] else "✗"
print(f" {marker} {name}: {info['value']}")
# 3. Проверка путей разморозки (для Phase B)
unfreeze = check_unfreeze_layers(student, UNFREEZE_LAYERS)
print(f"\nUnfreeze check (Phase B):")
for name, info in unfreeze.items():
marker = "✓" if info["found"] else "✗"
print(f" {marker} {name}: {info['n_params']:,} params "
f"({info['n_tensors']} tensors)")
# 4. Сохранение
OUT_MODEL.parent.mkdir(parents=True, exist_ok=True)
torch.save(student.state_dict(), OUT_MODEL)
size_mb = OUT_MODEL.stat().st_size / 1e6
print(f"\nSaved: {OUT_MODEL} ({size_mb:.1f} MB)")
AutoTokenizer.from_pretrained(ENCODER).save_pretrained(OUT_TOKENIZER)
print(f"Saved tokenizer: {OUT_TOKENIZER}")
# 5. Meta
meta = {
"script": "4.1_build_student",
"timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
"encoder": ENCODER,
"config": {
"hidden": HIDDEN,
"teacher_dim": TEACHER_DIM,
"num_heads": NUM_HEADS,
"seed": SEED,
},
"params": params,
"state_dict_size_mb": round(size_mb, 1),
"unfreeze_layers": UNFREEZE_LAYERS,
"unfreeze_check": unfreeze,
"sanity": sanity,
"versions": {
"torch": torch.__version__,
"transformers": transformers.__version__,
},
"elapsed_sec": round(time.time() - t0, 1),
}
META.write_text(json.dumps(meta, indent=2, ensure_ascii=False), encoding="utf-8")
print(f"Meta: {META}")
# 6. Warnings
if not all(info["found"] for info in unfreeze.values()):
print("\n[!] WARNING: не все слои для разморозки найдены! "
"Phase B не сможет их разморозить.")
if not all(info["ok"] for info in sanity.values()):
print("\n[!] WARNING: sanity checks провалились. "
"Проверьте конфиг энкодера.")
if __name__ == "__main__":
main()Ожидаемый вывод:
Loading encoder: sergeyzh/rubert-tiny-turbo
Loading weights: 100%|██████████████████████| 55/55 [00:00<00:00, 16624.87it/s]
Params:
encoder: 29,193,768
pooling: 1,248 (K=4)
projection: 240,384 (312→768, bias=True)
TOTAL: 29,435,400
Sanity checks:
✓ hidden_size: 312
✓ vocab_size: 83828
✓ max_position_embeddings: 2048
✓ num_hidden_layers: 3
Unfreeze check (Phase B):
✓ layer.1: 767,184 params (16 tensors)
✓ layer.2: 767,184 params (16 tensors)
Saved: /..../rubert-tiny2-distill/data/student_init.pt (117.8 MB)
Saved tokenizer: /..../rubert-tiny2-distill/data/student_tokenizer
Meta: /..../rubert-tiny2-distill/data/meta/4.1_build_student_meta.json4.5. Sanity-check перед обучением
Проверяет:
Forward работает в обоих режимах (312 / 768), нормы ≈ 1.
Маска корректна: padding не влияет на эмбеддинг (при разном кол-ве PAD).
Детерминизм: одинаковый вход → одинаковый выход.
Разные тексты → разные эмбеддинги (информационно).
Файл: scripts/4.2_check_student_init.py
"""Проверка StudentTiny до обучения.
Проверяет:
1. Forward работает в обоих режимах (312 / 768), нормы ≈ 1.
2. Маска корректна: padding не влияет на эмбеддинг (при разном кол-ве PAD).
3. Детерминизм: одинаковый вход → одинаковый выход.
4. Разные тексты → разные эмбеддинги (информационно).
Сохраняет meta/4.2_check_student_init_meta.json.
"""
import importlib.util
import json
import sys
import time
from pathlib import Path
import torch
from transformers import AutoModel, AutoTokenizer
ROOT = Path(__file__).resolve().parent.parent
DATA = ROOT / "data"
META_DIR = DATA / "meta"
META_DIR.mkdir(parents=True, exist_ok=True)
META = META_DIR / "4.2_check_student_init_meta.json"
SCRIPT_41 = Path(__file__).parent / "4.1_build_student.py"
# Динамическая загрузка: имя файла начинается с цифры, обычный import не сработает
_spec = importlib.util.spec_from_file_location("build_student", SCRIPT_41)
build_student = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(build_student)
StudentTiny = build_student.StudentTiny
ENCODER = build_student.ENCODER
HIDDEN = build_student.HIDDEN
TEACHER_DIM = build_student.TEACHER_DIM
OUT_MODEL = build_student.OUT_MODEL
OUT_TOKENIZER = build_student.OUT_TOKENIZER
DEVICE = torch.device("cpu")
NORM_TOL = 1e-4
def check_forward(student, tokenizer):
texts = [
"кошка сидит на стуле",
"кот уселся на стуле",
"самолёт летит в небе",
]
enc = tokenizer(texts, padding=True, truncation=True, max_length=64,
return_tensors="pt")
ids, mask = enc["input_ids"], enc["attention_mask"]
with torch.no_grad():
e312 = student(ids, mask, projected=False)
e768 = student(ids, mask, projected=True)
norms_312 = e312.norm(dim=-1).tolist()
norms_768 = e768.norm(dim=-1).tolist()
return {
"shape_312": list(e312.shape),
"shape_768": list(e768.shape),
"norms_312": [round(n, 6) for n in norms_312],
"norms_768": [round(n, 6) for n in norms_768],
"ok": (
list(e312.shape) == [3, HIDDEN]
and list(e768.shape) == [3, TEACHER_DIM]
and all(abs(n - 1.0) < NORM_TOL for n in norms_312)
and all(abs(n - 1.0) < NORM_TOL for n in norms_768)
),
}
def check_mask(student, tokenizer):
"""Padding не должен влиять на эмбеддинг, даже при разном кол-ве PAD."""
short = "кошка сидит на стуле"
long1 = "Это длинный текст, длиннее первого."
long2 = "Это ещё более длинный текст, длиннее предыдущего, чтобы short получил больше padding."
single = tokenizer([short], padding=True, return_tensors="pt")
batch1 = tokenizer([short, long1], padding=True, return_tensors="pt")
batch2 = tokenizer([short, long2], padding=True, return_tensors="pt")
# Плюс max_length padding — как в train_tiny.py
maxlen = tokenizer([short], padding="max_length", max_length=256,
truncation=True, return_tensors="pt")
with torch.no_grad():
e_single = student(single["input_ids"], single["attention_mask"], projected=False)
e_batch1 = student(batch1["input_ids"], batch1["attention_mask"], projected=False)
e_batch2 = student(batch2["input_ids"], batch2["attention_mask"], projected=False)
e_maxlen = student(maxlen["input_ids"], maxlen["attention_mask"], projected=False)
diffs = {
"batch1": float((e_single[0] - e_batch1[0]).abs().max()),
"batch2": float((e_single[0] - e_batch2[0]).abs().max()),
"max_length_256": float((e_single[0] - e_maxlen[0]).abs().max()),
}
return {
"lens": {
"single": int(single["input_ids"].shape[1]),
"batch1": int(batch1["input_ids"].shape[1]),
"batch2": int(batch2["input_ids"].shape[1]),
"maxlen": int(maxlen["input_ids"].shape[1]),
},
"diffs": {k: round(v, 8) for k, v in diffs.items()},
"ok": all(v < 1e-5 for v in diffs.values()),
}
def check_determinism(student, tokenizer):
text = "кошка сидит на стуле"
enc = tokenizer([text], padding=True, return_tensors="pt")
with torch.no_grad():
e1 = student(enc["input_ids"], enc["attention_mask"], projected=False)
e2 = student(enc["input_ids"], enc["attention_mask"], projected=False)
diff = float((e1 - e2).abs().max())
return {"diff": diff, "ok": diff < 1e-6}
def check_discrimination(student, tokenizer):
"""Информационно: до обучения семантика не гарантирована."""
texts = ["кошка сидит на стуле", "кот уселся на стуле", "самолёт летит в небе"]
enc = tokenizer(texts, padding=True, return_tensors="pt")
with torch.no_grad():
e = student(enc["input_ids"], enc["attention_mask"], projected=False)
return {
"cos_similar": round(float((e[0] * e[1]).sum()), 4),
"cos_different": round(float((e[0] * e[2]).sum()), 4),
"note": "до обучения семантика не гарантирована — informational",
}
def main():
t0 = time.time()
print(f"Device: {DEVICE}", flush=True)
tokenizer = AutoTokenizer.from_pretrained(OUT_TOKENIZER)
print(f"Loading encoder: {ENCODER}", flush=True)
encoder = AutoModel.from_pretrained(ENCODER)
student = StudentTiny(encoder).to(DEVICE)
sd = torch.load(OUT_MODEL, map_location=DEVICE)
if isinstance(sd, dict) and "state_dict" in sd:
sd = sd["state_dict"]
student.load_state_dict(sd)
student.eval()
n_total = sum(p.numel() for p in student.parameters())
print(f"Student loaded: {n_total:,} params\n")
checks = {}
print("[1] Forward (312 / 768), нормы ≈ 1:")
checks["forward"] = check_forward(student, tokenizer)
f = checks["forward"]
print(f" shapes: {f['shape_312']}, {f['shape_768']}")
print(f" norms 312: {f['norms_312']}")
print(f" norms 768: {f['norms_768']}")
print(f" {'✓' if f['ok'] else '✗'}")
print("\n[2] Mask sanity (padding не влияет):")
checks["mask"] = check_mask(student, tokenizer)
m = checks["mask"]
print(f" lens: {m['lens']}")
for k, v in m["diffs"].items():
print(f" diff[{k}]: {v:.8f}")
print(f" {'✓' if m['ok'] else '✗'}")
print("\n[3] Determinism:")
checks["determinism"] = check_determinism(student, tokenizer)
d = checks["determinism"]
print(f" diff: {d['diff']:.8f} {'✓' if d['ok'] else '✗'}")
print("\n[4] Discrimination (информационно):")
checks["discrimination"] = check_discrimination(student, tokenizer)
dc = checks["discrimination"]
print(f" cos(similar): {dc['cos_similar']}")
print(f" cos(different): {dc['cos_different']}")
# Только критические проверки влияют на all_ok
critical = ["forward", "mask", "determinism"]
all_ok = all(checks[k]["ok"] for k in critical)
meta = {
"script": "4.2_check_student_init",
"timestamp": time.strftime("%Y-%m-%dT%H:%M:%S"),
"n_params": n_total,
"norms_tolerance": NORM_TOL,
"checks": checks,
"critical_checks": critical,
"all_ok": all_ok,
"elapsed_sec": round(time.time() - t0, 2),
}
META.write_text(json.dumps(meta, indent=2, ensure_ascii=False), encoding="utf-8")
print(f"\nMeta: {META}")
if not all_ok:
print("\n[!] Критические проверки провалились.")
sys.exit(1)
if __name__ == "__main__":
main()Ожидаемый вывод:
Device: cpu
Loading encoder: sergeyzh/rubert-tiny-turbo
Loading weights: 100%|██████████████████████| 55/55 [00:00<00:00, 16529.57it/s]
Student loaded: 29,435,400 params
[1] Forward (312 / 768), нормы ≈ 1:
shapes: [3, 312], [3, 768]
norms 312: [1.0, 1.0, 1.0]
norms 768: [1.0, 1.0, 1.0]
✓
[2] Mask sanity (padding не влияет):
lens: {'single': 7, 'batch1': 9, 'batch2': 18, 'maxlen': 256}
diff[batch1]: 0.00000004
diff[batch2]: 0.00000004
diff[max_length_256]: 0.00000005
✓
[3] Determinism:
diff: 0.00000000 ✓
[4] Discrimination (информационно):
cos(similar): 0.8223
cos(different): 0.3381
Meta: /..../rubert-tiny2-distill/data/meta/4.2_check_student_init_meta.jsonЕсли Mask sanity: FAIL — pooling не маскирует padding. Это критично, потому что при обучении батчи формируются с padding'ом, и утечка через них испортит эмбеддинги.
Наблюдения:
discrimination: cos(similar)=0.8223, cos(different)=0.3381— уже до обучения случайный студент различает "кошка на стуле / кот уселся" (0.82) vs "кошка / самолёт" (0.34). Это заслуга предобученного энкодера turbo:hidden statesнесут семантику, а случайный attention pooling её не разрушает.determinism = 0.0— уже идеально.
KioskNews shows a cleaned-up reading view extracted from the publisher’s page — the original always lives on their site, not ours.