Часть 4. Создание ученика
Важное предупреждение:
Весь код по моей методологии написан нейросетью 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.json
4.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— уже идеально.
Ученик готов.