Looped LM 55.7M

Исследовательский decoder-only чекпойнт от KrzysztofZawisla, обученный с нуля на TinyStories (английские детские истории) с byte-level BPE на 16k. В архитектуре 2 prelude-блока, 2 общих core-блока применяются 8 раз и 2 coda-блока: уникальных transformer-блоков всего 6, но на инференсе модель считает 20 виртуальных слоёв. Это не production-ассистент, а research-модель.

👀 Ключевые особенности

• 🔥 Глубина задаётся циклом: shared core weights, sinusoidal loop encoding и input injection, который возвращает состояние к prelude на каждой итерации.
• 💪 n_loops — ручка вычислений на вызов: качество резко растёт на первых циклах и достигает пика на обученной глубине 8; дальше деградирует плавно.
• 🐱 4 цикла уже возвращают большую часть качества примерно за 63% вычислений; 1–2 цикла дают заметно деградированный текст.
• 🤫 Во время обучения на последних 2 циклах использовались auxiliary LM heads (logit lens).
• 👀 Это TinyStories-only модель: простой английский, короткие истории, без мировых знаний; нет KV cache, генерация пересчитывает весь контекст.

🐶 Пример

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
repo = "KrzysztofZawisla/looped-lm-tinystories-55M"
tok = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained(repo, trust_remote_code=True, dtype=torch.bfloat16)
inputs = tok("Once upon a time", return_tensors="pt")
out = model.generate(**inputs, max_new_tokens=150, n_loops=8, pad_token_id=tok.eos_token_id)
print(tok.decode(out[0], skip_special_tokens=True))

Подробнее

Валидация TinyStories: loss 1.2711, ppl 3.56 при 8 циклах; 4 цикла — loss 1.4605, ppl 4.31 при ~63% вычислений.
Обучение: 8000 шагов × 131,072 токена (~1.05B токенов), AdamW, cosine schedule, bf16.
Ограничения: только простой английский TinyStories, нет мировых знаний; нужен trust_remote_code=True, а llama.cpp / Ollama / vLLM без отдельной реализации не запустят.