Upload README.md with huggingface_hub

README.md

@@ -0,0 +1,393 @@
+---
+license: apache-2.0
+language:
+  - en
+library_name: pytorch
+tags:
+  - research
+  - transformer
+  - attention-residuals
+  - muon-optimizer
+  - nca-pretraining
+  - geometric-monitoring
+  - causal-lm
+datasets:
+  - allenai/peS2o
+  - open-web-math/open-web-math
+  - HuggingFaceTB/finemath
+  - bigcode/the-stack
+  - deepmind/pg19
+  - pile-of-law/pile-of-law
+  - OpenAssistant/oasst2
+pipeline_tag: text-generation
+model-index:
+  - name: kotodama-108m-base
+    results:
+      - task:
+          type: text-generation
+          name: Language Modeling
+        dataset:
+          type: wikitext
+          name: WikiText-2
+        metrics:
+          - name: Word Perplexity (fc-base)
+            type: perplexity
+            value: 41.76
+          - name: Word Perplexity (bcpt-base)
+            type: perplexity
+            value: 52.09
+      - task:
+          type: multiple-choice
+          name: ARC-Easy
+        dataset:
+          type: ai2_arc
+          name: ARC-Easy
+        metrics:
+          - name: Accuracy (fc-base)
+            type: accuracy
+            value: 0.455
+          - name: Accuracy (bcpt-base)
+            type: accuracy
+            value: 0.445
+---
+
+# Kotodama 108M Base
+
+A 108M parameter decoder-only transformer trained as a **proxy model** for validating architectural and optimizer choices before scaling to 3B parameters. This is a research artifact, not a production model.
+
+The model combines three techniques not previously studied together at this scale:
+
+- **Block Attention Residuals (AttnRes)** -- learned residual connections across transformer blocks that prevent BOS-sink attention collapse and produce 4x gradient uniformity across depth.
+- **NCA pre-pretraining** -- bootstrapping attention circuits using Neural Cellular Automata trajectories before language training, which trains attention patterns (not MLPs) and creates an L14 attractor basin in the representation manifold.
+- **Muon optimizer** -- spectral-norm steepest descent via Newton-Schulz orthogonalization, producing 2-4x higher stable rank than AdamW at matched loss, with Gram-NS optimized coefficients.
+
+**Organization:** [aethera-gp](https://huggingface.co/aethera-gp)
+**Training code:** [github.com/aethera-gp/kotodama](https://github.com/aethera-gp/kotodama) (pretraining/)
+
+## Architecture
+
+The model uses a Llama-family architecture with QK-norm and Block Attention Residuals.
+
+| Parameter | Value |
+|-----------|-------|
+| Parameters | 107.8M (+ 58.4K AttnRes) |
+| Hidden size | 512 |
+| Layers | 28 |
+| Query heads | 4 |
+| KV heads | 2 (GQA ratio 2:1) |
+| Head dim | 128 |
+| Intermediate size (SwiGLU) | 1408 |
+| Vocabulary | 49,152 (SmolLM2 tokenizer) |
+| Max context | 4,096 tokens |
+| Positional encoding | RoPE (theta=500,000) |
+| Normalization | Pre-RMSNorm + QK-norm |
+| Embeddings | Tied input/output |
+| Bias | None |
+| z-loss | 1e-5 |
+| AttnRes block boundaries | [0, 3, 7, 12, 21, 25] (DD-v1) |
+
+### Block Attention Residuals (DD-v1)
+
+AttnRes adds per-layer learned pseudo-queries and key norms that create residual connections between block boundaries. The DD-v1 configuration divides the 28-layer network into 6 variable-size blocks at layers [0, 3, 7, 12, 21, 25]. This adds only 58.4K parameters (0.05% overhead) but has substantial effects on training dynamics.
+
+Each transformer block stores:
+- `attn_res_query` / `attn_res_norm`: attention sub-block residual
+- `mlp_res_query` / `mlp_res_norm`: MLP sub-block residual
+
+A final `final_res_query` / `final_res_norm` aggregates block outputs before the LM head.
+
+### Differences from stock Llama
+
+- **QK-norm**: RMSNorm on Q and K projections after linear projection, enabling higher learning rates
+- **z-loss**: LSE-squared regularization preventing logit explosion
+- **Smaller vocab** (49K vs 128K): reduces the Godey gradient bottleneck (~94% destruction at 3072/49K vs ~98% at 3072/128K for the 3B target)
+- **Block AttnRes**: cross-block residual connections (see above)
+
+## Training
+
+### Optimizer Configuration
+
+Hybrid Muon + AdamW: Muon handles 2D weight matrices (Q/K/V/O projections, FFN gate/up/down -- ~77% of parameters), AdamW handles everything else (embeddings, norms).
+
+| Parameter | Muon (2D weights) | AdamW (embeddings, norms) |
+|-----------|-------------------|---------------------------|
+| Learning rate | 0.02 | 6e-4 |
+| Momentum / betas | 0.95 (Nesterov) | (0.9, 0.95) |
+| Weight decay | 0.01 | 0.1 |
+| NS iterations | 5 (Gram-NS coefficients) | -- |
+
+**Schedule:** WSD (Warmup-Stable-Decay). 5,000 step warmup (~6%), stable plateau to 90% of training, cosine decay over final 10%.
+
+**Gradient clipping:** 1.0
+
+**Precision:** BF16 autocast with FP8 compute (FP32 optimizer states).
+
+### NCA Pre-Pretraining
+
+Before language training, attention weights were bootstrapped using NCA (Neural Cellular Automata) pre-pretraining following Han et al. (2026). An NCA checkpoint co-trained with AttnRes DD-v1 (seed-17, 852M tokens) was used as initialization. After NCA, embeddings were reinitialized to the language vocabulary while attention weights, MLPs, and norms from NCA training were preserved (embed-only reinit).
+
+### Data Mix (Fullcorpus)
+
+170.4B tokens from 13 sources, shuffled with seed 42, sequence length 4096.
+
+| Source | Tokens | % | Category |
+|--------|--------|---|----------|
+| peS2o | 60.7B | 35.6% | Academic papers (Semantic Scholar) |
+| OpenCoderReasoning | 35.7B | 21.0% | Code reasoning (R1 + QwQ, Python/C++) |
+| Pile of Law | 18.8B | 11.0% | Legal (court opinions, congressional) |
+| StackExchange | 15.7B | 9.2% | Q&A (22 high-value sites) |
+| OpenWebMath | 14.1B | 8.2% | Math web pages |
+| FineMath | 10.8B | 6.4% | Quality-scored math (4+ score) |
+| PG-19 | 7.5B | 4.4% | Books (Project Gutenberg, 71K) |
+| Wikipedia | 5.0B | 3.0% | English Wikipedia |
+| SmolTalk | 0.9B | 0.6% | Synthetic multi-turn dialogue |
+| WildChat | 0.5B | 0.3% | Real user-GPT conversations |
+| SODA | 0.3B | 0.2% | Synthetic social dialogue |
+| Enron | 0.3B | 0.2% | Corporate email |
+| OASST2 | 0.01B | <0.1% | Human multi-turn conversations |
+
+**Category breakdown:** Academic/knowledge 38.6%, code reasoning 21.0%, math 14.6%, legal 11.0%, Q&A 9.2%, books 4.4%, conversation 1.1%.
+
+### Hardware and Compute
+
+- **Hardware:** 8x NVIDIA B200 (single node, NVLink)
+- **Parallelism:** DDP (DistributedDataParallel)
+- **Throughput:** ~1.96M tokens/sec average
+- **Micro batch size:** 16 per GPU
+- **Global batch size:** 2,097,152 tokens (16 * 4096 * 8 GPUs * gradient accumulation)
+- **torch.compile:** enabled (4x throughput vs eager)
+
+## Model Variants
+
+This repository contains two checkpoints from the same model lineage:
+
+### fc-base (fullcorpus)
+
+**File:** `fc-base.pt.zst`
+
+The primary pretraining run. 170.4B tokens over 81,252 steps on the full 13-source data mix described above. Initialized from NCA+AttnRes checkpoint (seed-17, 852M NCA tokens). WSD schedule with cosine decay in the final 10%.
+
+| Metric | Value |
+|--------|-------|
+| Final loss | 2.081 |
+| Min loss | 1.982 (step 80,200) |
+| Final perplexity | 8.01 |
+| Tokens seen | 170.4B |
+| Tokens/param ratio | ~1,581x |
+
+### bcpt-base (books-CPT)
+
+**File:** `bcpt-base.pt.zst`
+
+Continued pretraining of the fullcorpus model on 36.2B tokens of book data from three Common Pile sources not present in the original data mix. Resumed from fullcorpus step 72,000 (pre-decay, 151B tokens seen) with fresh optimizer state and a new WSD schedule (500-step warmup, 90% stable, 10% cosine decay).
+
+| Source | Tokens | % |
+|--------|--------|---|
+| Pre-1929 Books (Internet Archive/HathiTrust) | 19.1B | 52.8% |
+| Library of Congress | 14.0B | 38.7% |
+| DOAB (Open Access Books) | 3.1B | 8.6% |
+
+OCR quality filter applied: documents with >5% garbage characters dropped.
+
+| Metric | Value |
+|--------|-------|
+| Final loss | 2.342 |
+| Min loss | 2.230 (step 17,260) |
+| Final perplexity | 10.40 |
+| Additional tokens | 36.4B (17,337 steps) |
+| Total tokens seen | ~187.4B (resumed from step 72K / 151B tokens) |
+
+The higher loss/perplexity relative to fullcorpus reflects the domain shift to OCR book text, not regression. The books-CPT variant trades general benchmark performance for improved performance on literary and long-form text.
+
+## Evaluation
+
+### LM-Eval Benchmarks
+
+All benchmarks run zero-shot via lm-evaluation-harness.
+
+| Benchmark | Metric | fc-base | bcpt-base |
+|-----------|--------|---------|-----------|
+| ARC-Easy | acc | 0.455 | 0.445 |
+| ARC-Easy | acc_norm | 0.387 | 0.388 |
+| BoolQ | acc | 0.559 | 0.499 |
+| COPA | acc | 0.590 | 0.590 |
+| HellaSwag | acc | 0.277 | 0.280 |
+| HellaSwag | acc_norm | 0.297 | 0.295 |
+| LAMBADA | acc | 0.281 | 0.297 |
+| LAMBADA | ppl | 83.3 | 85.5 |
+| PIQA | acc | 0.577 | 0.588 |
+| PIQA | acc_norm | 0.569 | 0.571 |
+| SciQ | acc | 0.783 | 0.779 |
+| SciQ | acc_norm | 0.700 | 0.685 |
+| WikiText | word_ppl | 41.76 | 52.09 |
+| WikiText | bits/byte | 1.007 | 1.066 |
+| Winogrande | acc | 0.508 | 0.515 |
+
+**Notes:** These are proxy-scale (108M) results. Performance is expected at this scale -- the model was not designed to maximize benchmarks. The books-CPT variant shows slight improvements on commonsense/physical reasoning (PIQA, Winogrande, LAMBADA accuracy) and slight degradation on knowledge-heavy tasks (BoolQ, WikiText perplexity), consistent with the domain shift toward literary text.
+
+## Analysis Highlights
+
+The primary value of this model as a research artifact is the geometric monitoring data collected during training. The analysis packages in `fc-analysis/` and `bcpt-analysis/` contain activation geometry, concept geometry, and full metric histories.
+
+### Geometric Health (Final Checkpoint)
+
+Monitored at layers [0, 7, 14, 21, 27] throughout training.
+
+| Metric | Value | Interpretation |
+|--------|-------|----------------|
+| RankMe (embedding) | 440.5 | High effective dimensionality (out of 512) |
+| RankMe rebound ratio | 15.9x | Strong recovery from early collapse (min 27.7 at step 150) |
+| WeightWatcher alpha | 7.71 | Within Muon-healthy range (see notes) |
+| TwoNN intrinsic dim | 5.76 | Representation manifold dimensionality |
+| Dead units | 0.0% | No dead neurons at any monitored layer |
+
+### Stable Rank Profiles Across Depth
+
+Stable rank (effective rank of weight matrices) remains high across all layers throughout training, a signature of Muon's balanced spectral updates. Representative values from the final checkpoint (step 81,225):
+
+| Layer | Q proj | K proj | O proj | Gate proj | Down proj |
+|-------|--------|--------|--------|-----------|-----------|
+| 0 | 18.7 | 15.7 | 46.3 | 127.0 | 56.8 |
+| 7 | 42.5 | 40.0 | 87.9 | 76.8 | 140.4 |
+| 14 | 49.1 | 41.5 | 43.1 | 70.2 | 125.0 |
+| 21 | 39.4 | 30.0 | 67.9 | 62.9 | 49.2 |
+| 27 | 43.8 | 32.3 | 115.3 | 76.2 | 127.8 |
+
+Key observations:
+- **No low-rank collapse:** All weight matrices maintain high stable rank through 170B tokens. Under AdamW, these values would typically be 2-4x lower.
+- **Depth utilization:** Non-monotonic stable rank profile indicates all layers are actively contributing (not degenerating into near-identity transformations).
+- **Zero dead units:** No layer shows any dead neurons, even after extreme overtraining (1,581x tokens/parameter).
+
+### Attention Entropy Across Depth
+
+| Layer | Mean Entropy | Std | Interpretation |
+|-------|-------------|-----|----------------|
+| 0 | 6.13 | 0.43 | Broad attention (early feature mixing) |
+| 7 | 4.64 | 0.77 | Selective attention with variance |
+| 14 | 5.49 | 0.41 | Moderate selectivity |
+| 21 | 5.68 | 0.29 | Moderate, low variance |
+| 27 | 4.14 | 0.79 | Most selective (prediction heads) |
+
+This gradient -- broad at the bottom, selective at the top -- is the healthy pattern. Crucially, **the deep layers (L27) maintain diverse attention patterns** (std=0.79) rather than collapsing to BOS-sink. In baseline models without AttnRes, layers 21-27 develop 89-90% BOS attention concentration by this training stage.
+
+### Anisotropy Profile
+
+| Layer | Anisotropy |
+|-------|-----------|
+| 0 | 0.066 |
+| 7 | 0.452 |
+| 14 | 0.413 |
+| 21 | 0.148 |
+| 27 | 0.090 |
+
+The inverted-U anisotropy profile (low at edges, peaking at middle layers) indicates structured representational geometry rather than isotropy collapse or extreme anisotropy.
+
+### AttnRes Effects (from Proxy Phase Ablations)
+
+These findings come from the 5-run optimizer sweep at 6B tokens and the full 170B run:
+
+- **BOS-sink prevention:** Baseline models develop 89-90% BOS attention at deep layers by 6B tokens. DD-v1 AttnRes prevents this entirely, maintaining diverse attention patterns at all depths.
+- **4x gradient uniformity:** Gradient norm variance across layers is ~4x lower with AttnRes, enabling more uniform learning across depth.
+- **Full depth utilization:** Without AttnRes, deep layers tend toward near-identity transformations. With AttnRes, stable rank and attention entropy remain diverse at all depths.
+- **DD-v2 fragility:** Shifting even one block boundary (L12 to L14) produced 12/16 geometric metrics outside the range of all other configurations. Variable-size blocks cascade nonlinearly.
+
+### NCA Pre-Pretraining Effects
+
+- **Trains attention, not MLPs:** NCA pre-pretraining primarily structures attention weight matrices. MLP weights show minimal structured change, confirming that MLP reinit after NCA is correct.
+- **L14 attractor basin:** NCA creates a distinctive geometric signature at layer 14 that persists through full language training. This basin is present regardless of AttnRes configuration.
+- **Sub-additive with AttnRes:** NCA + AttnRes produces only +0.008 nats over the better of either alone, but preserves geometric properties from both techniques everywhere in the network.
+
+## Key Findings (Proxy Phase)
+
+1. **Muon lr=0.02 is the Pareto optimum** for 108M: matches AdamW final loss while maintaining 2-4x higher stable rank across all weight matrices.
+2. **torch.compile is the dominant throughput optimization**, providing 4x improvement. Liger kernels without FusedLinearCE hurt compile by 13%.
+3. **Extreme overtraining (1,581x tokens/param) does not cause geometric collapse** with Muon + AttnRes. Stable rank, attention entropy, and dead unit counts all remain healthy at 170B tokens.
+4. **WW alpha healthy range is higher for Muon than AdamW.** Alpha values of 7-8 are normal for Muon-trained models; do not apply AdamW-calibrated thresholds (which would flag these as unhealthy).
+
+## Usage
+
+The checkpoints are stored as compressed PyTorch state dicts (`.pt.zst`). To load:
+
+```python
+import torch
+import zstandard as zstd
+import io
+
+# Decompress
+with open("fc-base.pt.zst", "rb") as f:
+    dctx = zstd.ZstdDecompressor()
+    decompressed = dctx.decompress(f.read())
+
+# Load state dict
+state_dict = torch.load(io.BytesIO(decompressed), map_location="cpu", weights_only=True)
+
+# Initialize model (requires the kotodama training code)
+from src.model.llama import LuxiaBaseModel, LuxiaModelConfig
+
+config = LuxiaModelConfig(
+    hidden_size=512,
+    num_layers=28,
+    num_attention_heads=4,
+    num_kv_heads=2,
+    head_dim=128,
+    intermediate_size=1408,
+    vocab_size=49152,
+    max_position_embeddings=4096,
+    rope_theta=500000.0,
+    qk_norm=True,
+    tie_word_embeddings=True,
+    z_loss_weight=1e-5,
+    attn_res=True,
+    attn_res_boundaries=[0, 3, 7, 12, 21, 25],
+)
+
+model = LuxiaBaseModel(config)
+model.load_state_dict(state_dict)
+```
+
+**Tokenizer:** `HuggingFaceTB/SmolLM2-135M` (49,152 vocab, byte-fallback).
+
+## Repository Contents
+
+```
+fc-base.pt.zst              # Fullcorpus final checkpoint (81,252 steps, 170.4B tokens)
+bcpt-base.pt.zst             # Books-CPT checkpoint (17,337 additional steps, 36.4B tokens)
+fc-analysis/                 # Fullcorpus analysis package
+  activation_geometry/       # Per-layer activation extractions
+  concept_geometry/          # Concept-level geometric analysis
+  lm_eval/                   # Full lm-evaluation-harness results
+  report.html                # Analysis report
+bcpt-analysis/               # Books-CPT analysis package (same structure)
+fc-metrics.jsonl             # Fullcorpus training metrics (loss, LR, throughput)
+fc-geo_metrics.jsonl         # Fullcorpus geometric monitoring (stable rank, entropy, etc.)
+bcpt-metrics.jsonl           # Books-CPT training metrics
+bcpt-geo_metrics.jsonl       # Books-CPT geometric monitoring
+```
+
+## Limitations
+
+- **108M proxy scale.** This model exists to validate architecture and optimizer choices, not to be useful for downstream tasks. Benchmark performance reflects this.
+- **No raw code in training data.** The 645GB cleaned stack_v1 JSONL (~126B tokens, 130 languages) was never tokenized and is absent from the data mix. The model sees code only through reasoning traces (OpenCoderReasoning) and Q&A (StackExchange).
+- **Conversational data < 1.2%.** The original spec targeted 25% conversational data. The actual mix is dominated by academic text (35.6%) and code reasoning (21.0%).
+- **OCR noise in books-CPT.** Despite filtering documents with >5% garbage characters, the books-CPT data (pre-1929 scans, Library of Congress) contains residual OCR artifacts.
+- **No deduplication** was applied to the books-CPT data (estimated minimal cross-source overlap between digitization projects, but not verified).
+- **Eval methodology:** Top-p sampling catastrophically degrades generation quality at 108M scale. All evaluation uses pure temperature sampling only.
+
+## Citation
+
+```bibtex
+@misc{kotodama2026,
+  title={Kotodama: Block Attention Residuals and NCA Pre-Pretraining for Transformer Language Models},
+  author={Aethera GP},
+  year={2026},
+  url={https://huggingface.co/aethera-gp/kotodama-108m-base}
+}
+```
+
+### References
+
+- Block Attention Residuals: see `Attention_Residuals.pdf` in the training repo
+- NCA Pre-Pretraining: [Han et al., 2026](https://arxiv.org/abs/2603.10055)
+- Muon Optimizer: [MoonshotAI/Muon](https://github.com/MoonshotAI/Muon); [Moonlight: Muon is Scalable for LLM Training](https://arxiv.org/abs/2502.16982)
+- Gram-Newton-Schulz: [Dao-AILab/Gram-Newton-Schulz](https://github.com/Dao-AILab/Gram-Newton-Schulz)
+- WeightWatcher: [Martin et al.](https://arxiv.org/abs/2102.11258)