Step 1000: Int 4.59, Avg 34.55%, BPB 1.6713

README.md

@@ -0,0 +1,186 @@
+---
+library_name: transformers
+pipeline_tag: text-generation
+tags:
+- custom_code
+- causal-lm
+- text-generation
+- pytorch
+- ac-swiglu
+- inclusive-mixing
+- small-language-model
+- generalist
+- 4k-tokenizer
+datasets:
+- HuggingFaceFW/fineweb_edu_100BT-shuffled
+- HuggingFaceTB/smollm-corpus
+- epfml/FineWeb-HQ
+- HuggingFaceTB/finemath
+---
+
+# Inclusive AC-SwiGLU Mini
+
+Evaluated training checkpoint from a 4.94M-parameter
+inclusive attention-mixed SwiGLU
+generalist language model using the third-party
+AxiomicLabs GPT-S 4,096-token tokenizer. It has no place embeddings, role
+embeddings, or inference-time equation detection. It was recorded at step
+1,000 with WikiText normalized BPB 1.6713.
+
+## Loading
+
+This is a custom Transformers architecture. `trust_remote_code=True` is
+required because stock Hugging Face model classes do not implement AC-SwiGLU or this
+model's exact rotary convention.
+
+```bash
+pip install "torch>=2.5" "transformers>=4.50" safetensors
+```
+
+```python
+import torch
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+repo = "User01110/attention-contraction-mini"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+dtype = torch.bfloat16 if device == "cuda" else torch.float32
+tokenizer = AutoTokenizer.from_pretrained(repo, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    repo,
+    trust_remote_code=True,
+    torch_dtype=dtype,
+).to(device).eval()
+
+inputs = tokenizer("The process of photosynthesis", return_tensors="pt").to(device)
+output = model.generate(**inputs, max_new_tokens=64)
+print(tokenizer.decode(output[0], skip_special_tokens=True))
+```
+
+Checkpoint tensors are stored in bfloat16. Pass `torch_dtype=torch.float32`
+when an FP32 runtime is required; every stored BF16 value widens exactly to
+FP32, though the pre-export FP32 master-weight mantissa cannot be reconstructed.
+The model intentionally sets `use_cache=False`: generation is standard
+Transformers generation, but the visible context is recomputed for every new
+token because this experimental architecture does not implement a KV cache.
+
+## Architecture
+
+- Parameters: 4,943,712, with tied input/output embeddings
+- Weights: native bfloat16 safetensors (`model.safetensors`); no `.bin` weights
+- Runtime: native SDPA, a fused SwiGLU input projection, and small dense native
+  PyTorch expanded-space attention-mixing matmuls
+- Tokenizer: `AxiomicLabs/GPT-S-5M` at revision
+  `275b9c3ca78736bf6aeb154c7e2d5f5764fe9035`
+- Vocabulary: 4,096 third-party tokens
+- Parameter allocation: 884,736 tied embedding parameters and
+  4,058,976 non-embedding parameters
+- Supported/exported context: 1,024 tokens
+- Training block length: 1,024 tokens
+- Standalone prompt tokenization automatically prepends the native BOS token
+- Width/layers: 216 / 10
+- Token-attention heads: 6 query, 2 KV
+- Native SwiGLU expands each token to 432 channels with
+  one fused projection, applies `up * SiLU(gate)`, and retains an ordinary
+  zero-initialized linear contraction into the residual stream
+- Inclusive AC mixing uses all
+  18 expanded chunks as queries, keys,
+  and values, each of 24 channels split over 3 heads
+- Every head computes an exact dense
+  18x18
+  softmax mixing graph; no edge is removed or approximated
+- The inclusive operator is `V_mix = [I + s(A(x) - A0)]V`; `V_mix` then passes
+  through the ordinary SwiGLU down projection exactly once
+- There is no parallel AC output branch. At neutral mixing `A(x) = A0`, the
+  architecture is exactly native SwiGLU with `V_mix = V`
+- The learned key scale and zero-initialized down projection make the complete
+  AC-SwiGLU residual branch exactly zero at initialization without a route gate
+- The difference of two row-stochastic mixtures has zero row sum and bounded
+  L1 norm, providing a stable internal deformation of the expanded features
+- Contiguous-half RoPE without scaling
+- No task-specific model features or inference-time benchmark handling
+
+## Intended use and limitations
+
+This is a small base language model released for architecture research,
+representation analysis, and controlled comparisons. It is not instruction
+tuned and should not be treated as a factual authority or used for consequential
+decisions. Its 4.94M-parameter scale, English-weighted web mixture,
+1,024-token context, and cache-free generation materially limit
+capability and throughput. The usual web-corpus biases and inaccuracies remain.
+
+## Tokenizer provenance
+
+The tokenizer and its 4,096-token vocabulary were **not
+created or owned by the AC-SwiGLU model author**. They are reused from the public
+[`AxiomicLabs/GPT-S-5M`](https://huggingface.co/AxiomicLabs/GPT-S-5M)
+repository at the exact revision listed above, whose repository metadata
+identifies Axiomic Labs as the publisher and Apache-2.0 as the license. No
+claim of tokenizer ownership beyond that public attribution is made here. The
+exported copy preserves its vocabulary and tokenization pipeline; AC-SwiGLU only
+configures its existing `<bos>` token to be prepended automatically and raises
+the exported model/tokenizer context limit to 1,024 tokens.
+
+## Optimization
+
+- Training budget: 20,971,520,000 tokens over 40,000 updates
+- Effective batch: 524,288 tokens per update
+- Native microbatch: 512 sequences x 1,024 tokens,
+  accumulated 1 times
+- Runtime implementation: explicit BF16 residual/AC-SwiGLU activations, one
+  fused expansion GEMM, dense 18x18 inclusive mixing matmuls, and full-model
+  `torch.compile` enabled by default
+- Runtime: 1 resumable 40,000-update sessions
+  with exact full-precision optimizer/RNG/data-stream recovery checkpoints
+- Learning rate: 1,000-update linear warmup, flat at 2.5e-03 through update 30,000, then cosine to zero on update 40,000; original-scaling Muon follows the same
+  multiplier from its 3.0e-02 peak
+- The final configured update has exactly zero learning rate
+- Official PyTorch Muon with original matrix-shape scaling for hidden matrices;
+  AdamW for embeddings and remaining parameters
+
+## Training mixture
+
+- FineWeb-Edu 100BT shuffled: 55%
+- Cosmopedia v2: 25%
+- FineWeb-HQ: 10%
+- FineMath 4+: 10%
+
+FineWeb-Edu supplies the primary educational web text, Cosmopedia supplies
+synthetic textbook-style coverage, FineWeb-HQ contributes model-filtered,
+knowledge-rich general web text, and FineMath-4+ supplies mathematical
+reasoning as ordinary causal-language-model text. The mixture remains fixed
+for the complete run.
+There are no benchmark labels or benchmark-specific preprocessing. Dataset
+revisions are pinned and the same fixed mixture is maintained across both
+runtime sessions.
+
+## Zero-shot evaluation at step 1,000
+
+The four lm-eval tasks use normalized accuracy when supplied by lm-eval
+0.4.12, with native bfloat16 weights and float32
+likelihood softmax. ArithMark-3 uses its official primary `acc_norm` metric:
+mean continuation-token log likelihood, with context and continuation tokenized
+separately after one native BOS prefix. Autocast is not used for evaluation. The four
+lm-eval benchmark contexts use the tokenizer's native BOS behavior.
+
+**Open SLM Intelligence Index: 4.59**
+**Open SLM Average: 34.55%**
+
+| Benchmark | Accuracy |
+|---|---:|
+| HellaSwag | 27.49% |
+| ARC-Easy | 31.14% |
+| ARC-Challenge | 20.05% |
+| PIQA | 54.19% |
+| ArithMark-3 (`acc_norm`) | 29.70% |
+
+The Average weights HellaSwag, combined ARC (the mean of ARC-Easy and
+ARC-Challenge), and PIQA at 1.0, and ArithMark-3 at
+0.75. The Intelligence Index first maps random
+chance to 0 and perfect accuracy to 100, then applies those same weights,
+matching Open SLM Leaderboard revision
+`2fbaaa164009c6c0d201ddf245878f5942cc8628`.
+
+WikiText-103 validation at this step: loss 3.6964, perplexity
+40.30, normalized BPB 1.6713 over 359,037
+scored tokens and 1,145,591 normalized UTF-8 bytes, using one initial BOS,
+1,024-token windows, and a 512-token stride.

benchmark_results.json

@@ -0,0 +1,202 @@
+{
+  "step": 1000,
+  "parameters": 4943712,
+  "lm_eval_version": "0.4.12",
+  "evaluation_dtype": "bfloat16",
+  "softmax_dtype": "float32",
+  "evaluation_autocast": false,
+  "num_fewshot": 0,
+  "model_context": 1024,
+  "training_context": 1024,
+  "lm_eval_context": 1024,
+  "lm_eval_bos_prefix": true,
+  "arithmark_3.0_bos_prefix": true,
+  "benchmark_accuracy": {
+    "arc_easy": 0.3114478114478115,
+    "arc_challenge": 0.20051194539249148,
+    "hellaswag": 0.2749452300338578,
+    "piqa": 0.5418933623503809
+  },
+  "lm_eval_metric_policy": "acc_norm when available, otherwise acc",
+  "intelligence_index": 4.587205403718424,
+  "average_accuracy_percent": 34.548492554783735,
+  "open_slm_scores": {
+    "leaderboard": "AxiomicLabs/Open_SLM_Leaderboard",
+    "leaderboard_revision": "2fbaaa164009c6c0d201ddf245878f5942cc8628",
+    "intelligence_index": 4.587205403718424,
+    "average_accuracy_percent": 34.548492554783735,
+    "components_percent": {
+      "hellaswag": 27.49452300338578,
+      "combined_arc": 25.59798784201515,
+      "piqa": 54.18933623503809,
+      "arithmark_3.0": 29.7
+    },
+    "chance_normalized_components": {
+      "hellaswag": 3.3260306711810395,
+      "combined_arc": 0.7973171226868677,
+      "piqa": 8.378672470076182,
+      "arithmark_3.0": 6.266666666666666
+    },
+    "average_accuracy_weights": {
+      "hellaswag": 1.0,
+      "combined_arc": 1.0,
+      "piqa": 1.0,
+      "arithmark_3.0": 0.75
+    },
+    "intelligence_index_weights": {
+      "hellaswag": 1.0,
+      "combined_arc": 1.0,
+      "piqa": 1.0,
+      "arithmark_3.0": 0.75
+    }
+  },
+  "checkpoint_policy": "every completed evaluation is committed regardless of whether validation BPB is best; rolling full-precision model, optimizer, RNG, stream cursor, and prefetched batches are stored separately",
+  "best_tracking_metric": "validation.normalized_bpb@1024 (lower is better)",
+  "best_validation": {
+    "normalized_bpb": 1.6713139089025868,
+    "step": 1000
+  },
+  "arithmark_3.0": {
+    "dataset": "AxiomicLabs/Arithmark-3.0",
+    "dataset_revision": "6f6e59dd9b7e2c63455f7af7f838f9ecc3d0a746",
+    "primary_metric": "acc_norm",
+    "metric_unit": "fraction",
+    "acc": 0.314,
+    "acc_norm": 0.297,
+    "raw_correct": 314,
+    "normalized_correct": 297,
+    "total": 1000,
+    "max_context": 1024,
+    "by_category": {
+      "elementary_school_math_continuation::addition::grades_1_2::easy": {
+        "acc": 0.21875,
+        "acc_norm": 0.2109375,
+        "raw_correct": 28,
+        "normalized_correct": 27,
+        "total": 128
+      },
+      "elementary_school_math_continuation::comparison::grades_2_3::medium": {
+        "acc": 0.1590909090909091,
+        "acc_norm": 0.13636363636363635,
+        "raw_correct": 7,
+        "normalized_correct": 6,
+        "total": 44
+      },
+      "elementary_school_math_continuation::comparison_difference::grades_2_3::medium": {
+        "acc": 0.25,
+        "acc_norm": 0.25,
+        "raw_correct": 12,
+        "normalized_correct": 12,
+        "total": 48
+      },
+      "elementary_school_math_continuation::data::grades_2_3::easy": {
+        "acc": 0.16279069767441862,
+        "acc_norm": 0.16279069767441862,
+        "raw_correct": 7,
+        "normalized_correct": 7,
+        "total": 43
+      },
+      "elementary_school_math_continuation::division::grades_3_4::medium": {
+        "acc": 0.2962962962962963,
+        "acc_norm": 0.2962962962962963,
+        "raw_correct": 16,
+        "normalized_correct": 16,
+        "total": 54
+      },
+      "elementary_school_math_continuation::fractions_counting::grades_3_4::medium": {
+        "acc": 0.22,
+        "acc_norm": 0.22,
+        "raw_correct": 11,
+        "normalized_correct": 11,
+        "total": 50
+      },
+      "elementary_school_math_continuation::geometry_area::grades_4_5::medium": {
+        "acc": 0.4423076923076923,
+        "acc_norm": 0.34615384615384615,
+        "raw_correct": 23,
+        "normalized_correct": 18,
+        "total": 52
+      },
+      "elementary_school_math_continuation::geometry_perimeter::grades_4_5::medium": {
+        "acc": 0.6222222222222222,
+        "acc_norm": 0.5777777777777777,
+        "raw_correct": 28,
+        "normalized_correct": 26,
+        "total": 45
+      },
+      "elementary_school_math_continuation::measurement::grades_2_3::easy": {
+        "acc": 0.3157894736842105,
+        "acc_norm": 0.2631578947368421,
+        "raw_correct": 24,
+        "normalized_correct": 20,
+        "total": 76
+      },
+      "elementary_school_math_continuation::money::grades_3_4::medium": {
+        "acc": 0.265625,
+        "acc_norm": 0.25,
+        "raw_correct": 17,
+        "normalized_correct": 16,
+        "total": 64
+      },
+      "elementary_school_math_continuation::multiplication::grades_3_4::medium": {
+        "acc": 0.35135135135135137,
+        "acc_norm": 0.32432432432432434,
+        "raw_correct": 26,
+        "normalized_correct": 24,
+        "total": 74
+      },
+      "elementary_school_math_continuation::patterns::grades_3_4::medium": {
+        "acc": 0.3018867924528302,
+        "acc_norm": 0.2830188679245283,
+        "raw_correct": 16,
+        "normalized_correct": 15,
+        "total": 53
+      },
+      "elementary_school_math_continuation::subtraction::grades_1_2::easy": {
+        "acc": 0.21367521367521367,
+        "acc_norm": 0.2222222222222222,
+        "raw_correct": 25,
+        "normalized_correct": 26,
+        "total": 117
+      },
+      "elementary_school_math_continuation::time::grades_2_3::easy": {
+        "acc": 0.9454545454545454,
+        "acc_norm": 0.8909090909090909,
+        "raw_correct": 52,
+        "normalized_correct": 49,
+        "total": 55
+      },
+      "elementary_school_math_continuation::two_step_add_subtract::grades_2_3::medium": {
+        "acc": 0.1956521739130435,
+        "acc_norm": 0.2391304347826087,
+        "raw_correct": 9,
+        "normalized_correct": 11,
+        "total": 46
+      },
+      "elementary_school_math_continuation::two_step_addition::grades_2_3::medium": {
+        "acc": 0.10526315789473684,
+        "acc_norm": 0.10526315789473684,
+        "raw_correct": 2,
+        "normalized_correct": 2,
+        "total": 19
+      },
+      "elementary_school_math_continuation::two_step_subtraction::grades_2_3::medium": {
+        "acc": 0.34375,
+        "acc_norm": 0.34375,
+        "raw_correct": 11,
+        "normalized_correct": 11,
+        "total": 32
+      }
+    },
+    "choice_batch_size": 64
+  },
+  "validation": {
+    "loss": 3.696356708225175,
+    "perplexity": 40.300211143236396,
+    "normalized_bpb": 1.6713139089025868,
+    "tokens": 359037,
+    "normalized_utf8_bytes": 1145591,
+    "window": 1024,
+    "stride": 512
+  }
+}

config.json

@@ -0,0 +1,39 @@
+{
+  "architectures": [
+    "ACSwiGLUForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "modeling_ac_swiglu.ACSwiGLUConfig",
+    "AutoModelForCausalLM": "modeling_ac_swiglu.ACSwiGLUForCausalLM"
+  },
+  "model_type": "ac_swiglu",
+  "vocab_size": 4096,
+  "seq_len": 1024,
+  "max_position_embeddings": 1024,
+  "n_positions": 1024,
+  "n_ctx": 1024,
+  "d_model": 216,
+  "n_layers": 10,
+  "n_heads": 6,
+  "n_kv_heads": 2,
+  "chunk": 24,
+  "ac_heads": 3,
+  "expand": 2,
+  "ac_mix_scale": 1.0,
+  "mix_mode": "inclusive_centered_expanded_self_attention",
+  "ac_conditioning": "swiglu_hidden_queries_keys_values",
+  "ac_mixing_backend": "dense_18_by_18_softmax",
+  "tokenizer_source": "AxiomicLabs/GPT-S-5M",
+  "tokenizer_revision": "275b9c3ca78736bf6aeb154c7e2d5f5764fe9035",
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "pad_token_id": 2,
+  "tie_word_embeddings": true,
+  "is_decoder": true,
+  "is_encoder_decoder": false,
+  "use_cache": false,
+  "safe_serialization": true,
+  "minimum_torch_version": "2.5",
+  "torch_dtype": "bfloat16",
+  "transformers_version": "5.14.1"
+}

generation_config.json

@@ -0,0 +1,8 @@
+{
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "pad_token_id": 2,
+  "do_sample": false,
+  "use_cache": false,
+  "repetition_penalty": 1.2
+}

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f4541226faf35b590a94cf426b0de2e387df5bf1a5d718bbb7951a5c57923de0
+size 9973768

modeling_ac_swiglu.py

@@ -0,0 +1,415 @@
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+try:
+    from transformers import GenerationMixin
+except ImportError:
+    from transformers.generation import GenerationMixin
+from transformers import PretrainedConfig, PreTrainedModel
+from transformers.modeling_outputs import CausalLMOutputWithPast
+
+
+class ACSwiGLUConfig(PretrainedConfig):
+    model_type = "ac_swiglu"
+    attribute_map = {
+        "hidden_size": "d_model",
+        "num_hidden_layers": "n_layers",
+        "num_attention_heads": "n_heads",
+        "num_key_value_heads": "n_kv_heads",
+    }
+
+    def __init__(
+        self,
+        vocab_size=4096,
+        seq_len=1024,
+        d_model=216,
+        n_layers=10,
+        n_heads=6,
+        n_kv_heads=2,
+        chunk=24,
+        ac_heads=3,
+        expand=2,
+        ac_mix_scale=1.0,
+        max_position_embeddings=None,
+        n_positions=None,
+        n_ctx=None,
+        tie_word_embeddings=True,
+        **kwargs,
+    ):
+        if min(vocab_size, seq_len, d_model, n_layers) <= 0:
+            raise ValueError("Vocabulary, context, width, and depth must be positive.")
+        if min(n_heads, n_kv_heads, chunk, ac_heads, expand) <= 0:
+            raise ValueError(
+                "Attention, AC-SwiGLU, and expansion dimensions must be positive."
+            )
+        if d_model % n_heads or n_heads % n_kv_heads:
+            raise ValueError(
+                "d_model % n_heads and n_heads % n_kv_heads must be zero."
+            )
+        if (d_model // n_heads) % 2:
+            raise ValueError("The token-attention head dimension must be even for RoPE.")
+        if d_model % chunk or (d_model * expand) % chunk or chunk % ac_heads:
+            raise ValueError(
+                "AC-SwiGLU requires d_model and d_model * expand divisible by "
+                "chunk, and chunk divisible by ac_heads."
+            )
+        if not math.isfinite(ac_mix_scale) or ac_mix_scale <= 0.0:
+            raise ValueError("ac_mix_scale must be finite and positive.")
+        kwargs.setdefault("is_decoder", True)
+        kwargs.setdefault("is_encoder_decoder", False)
+        kwargs.setdefault("tie_word_embeddings", tie_word_embeddings)
+        # AC-SwiGLU exports intentionally recompute the full visible context during
+        # generation; this architecture does not expose a KV cache.
+        kwargs.setdefault("use_cache", False)
+        super().__init__(**kwargs)
+        self.vocab_size = vocab_size
+        self.seq_len = seq_len
+        self.max_position_embeddings = max_position_embeddings or seq_len
+        self.n_positions = n_positions or self.max_position_embeddings
+        self.n_ctx = n_ctx or self.max_position_embeddings
+        self.d_model = d_model
+        self.n_layers = n_layers
+        self.n_heads = n_heads
+        self.n_kv_heads = n_kv_heads
+        self.chunk = chunk
+        self.ac_heads = ac_heads
+        self.expand = expand
+        self.ac_mix_scale = ac_mix_scale
+        self.head_dim = d_model // n_heads
+        self.num_key_value_groups = n_heads // n_kv_heads
+        self.is_decoder = True
+        self.use_cache = False
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, d):
+        super().__init__()
+        self.w = nn.Parameter(torch.ones(d))
+
+    def forward(self, x):
+        return F.rms_norm(
+            x, (x.shape[-1],), self.w.to(dtype=x.dtype), eps=1e-6
+        )
+
+
+def rope_cache(seq_len, head_dim, device, base=10000.0):
+    inv = 1.0 / (
+        base ** (torch.arange(0, head_dim, 2, device=device).float() / head_dim)
+    )
+    t = torch.arange(seq_len, device=device).float()
+    freqs = torch.outer(t, inv)
+    return torch.cos(freqs), torch.sin(freqs)
+
+
+def apply_rope(x, cos, sin):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], dim=-1)
+
+
+class TokenAttention(nn.Module):
+    def __init__(self, d, n_heads, n_kv_heads):
+        super().__init__()
+        if d % n_heads or n_heads % n_kv_heads:
+            raise ValueError("d_model % n_heads and n_heads % n_kv_heads must be zero.")
+        self.h, self.kv_h, self.hd = n_heads, n_kv_heads, d // n_heads
+        version = torch.__version__.split("+", 1)[0].split(".")
+        self.native_gqa = tuple(map(int, version[:2])) >= (2, 5)
+        self.q = nn.Linear(d, d, bias=False)
+        self.k = nn.Linear(d, n_kv_heads * self.hd, bias=False)
+        self.v = nn.Linear(d, n_kv_heads * self.hd, bias=False)
+        self.o = nn.Linear(d, d, bias=False)
+        self.qn, self.kn = RMSNorm(self.hd), RMSNorm(self.hd)
+
+    def forward(self, x, cos, sin):
+        B, T, d = x.shape
+        q = self.q(x).view(B, T, self.h, self.hd).transpose(1, 2)
+        k = self.k(x).view(B, T, self.kv_h, self.hd).transpose(1, 2)
+        v = self.v(x).view(B, T, self.kv_h, self.hd).transpose(1, 2)
+        q, k = self.qn(q), self.kn(k)
+        q, k = apply_rope(q, cos, sin), apply_rope(k, cos, sin)
+        if self.h != self.kv_h and not self.native_gqa:
+            repeats = self.h // self.kv_h
+            k = k.repeat_interleave(repeats, dim=1)
+            v = v.repeat_interleave(repeats, dim=1)
+            y = F.scaled_dot_product_attention(q, k, v, is_causal=True)
+        else:
+            y = F.scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                is_causal=True,
+                enable_gqa=self.h != self.kv_h,
+            )
+        return self.o(y.transpose(1, 2).reshape(B, T, d))
+
+
+class ACSwiGLU(nn.Module):
+    # Fused SwiGLU with centered attention inside its expanded representation.
+
+    def __init__(self, d, chunk=24, heads=3, expand=2, mix_scale=1.0):
+        super().__init__()
+        expanded = d * expand
+        if min(d, chunk, heads, expand) <= 0:
+            raise ValueError("AC-SwiGLU dimensions and expansion must be positive.")
+        if d % chunk or expanded % chunk or chunk % heads:
+            raise ValueError(
+                "AC-SwiGLU requires d and d * expand divisible by chunk, "
+                "and chunk divisible by heads."
+            )
+        if not math.isfinite(mix_scale) or mix_scale <= 0.0:
+            raise ValueError("ac_mix_scale must be finite and positive.")
+        self.d, self.m, self.c, self.h = d, expanded, chunk, heads
+        self.n = expanded // chunk
+        self.hd, self.expand = chunk // heads, expand
+        self.mix_scale = float(mix_scale)
+        self.in_proj = nn.Linear(d, 2 * expanded, bias=False)
+        self.down = nn.Linear(expanded, d, bias=False)
+        self.q_scale = nn.Parameter(torch.ones(chunk))
+        self.k_scale = nn.Parameter(torch.zeros(chunk))
+        self.bias = nn.Parameter(torch.zeros(heads, self.n, self.n))
+        nn.init.zeros_(self.down.weight)
+
+    def _mix_expanded(self, x):
+        B, T, _ = x.shape
+        batch_tokens = B * T
+        up, gate = self.in_proj(x).chunk(2, dim=-1)
+        hidden = up * F.silu(gate)
+        values = hidden.reshape(
+            batch_tokens, self.n, self.h, self.hd
+        ).transpose(1, 2)
+        normalized = F.rms_norm(values, (self.hd,), eps=1e-6)
+        q = normalized * self.q_scale.view(1, self.h, 1, self.hd).to(
+            normalized.dtype
+        )
+        k = normalized * self.k_scale.view(1, self.h, 1, self.hd).to(
+            normalized.dtype
+        )
+        logits = (q @ k.transpose(-2, -1)) * (self.hd ** -0.5)
+        logits = logits + self.bias.unsqueeze(0).to(logits.dtype)
+        attn = F.softmax(logits, dim=-1, dtype=torch.float32).to(values.dtype)
+        reference = F.softmax(
+            self.bias.float(), dim=-1
+        ).to(values.dtype).unsqueeze(0)
+        delta_attn = attn - reference
+        mixed = values + self.mix_scale * (delta_attn @ values)
+        return hidden, values, mixed, logits, attn, delta_attn
+
+    def forward(self, x):
+        B, T, _ = x.shape
+        _, _, mixed, _, _, _ = self._mix_expanded(x)
+        mixed = mixed.transpose(1, 2).reshape(B, T, self.m)
+        return self.down(mixed)
+
+    @torch.no_grad()
+    def diagnostic_stats(self, x):
+        probe_count = x.shape[0]
+        B, T, _ = x.shape
+        hidden, values, mixed, logits, attn, delta_attn = self._mix_expanded(x)
+        mixed_hidden = mixed.transpose(1, 2).reshape(B, T, self.m)
+        base_output = self.down(hidden)
+        output = self.down(mixed_hidden)
+        flat_output = output.reshape(-1, self.d)
+        token_effect = (
+            flat_output.float().norm(dim=-1)
+            / x.reshape(-1, self.d).float().norm(dim=-1).clamp_min(1e-12)
+        )
+        interaction_effect = (
+            (output - base_output).reshape(-1, self.d).float().norm(dim=-1)
+            / flat_output.float().norm(dim=-1).clamp_min(1e-12)
+        )
+        probs = attn.float().clamp_min(1e-9)
+        entropy = -(probs * probs.log()).sum(dim=-1) / math.log(self.n)
+        identity_coefficient = (
+            1.0
+            + self.mix_scale
+            * delta_attn.float().diagonal(dim1=-2, dim2=-1)
+        ).mean()
+        if self.h > 1:
+            flattened = probs.transpose(0, 1).reshape(self.h, -1)
+            normalized = F.normalize(flattened, dim=-1)
+            similarity = normalized @ normalized.mT
+            head_similarity = (
+                similarity.sum() - similarity.diagonal().sum()
+            ) / (self.h * (self.h - 1))
+        else:
+            head_similarity = probs.new_zeros(())
+        mix_delta = mixed - values
+        retention_cosine = F.cosine_similarity(
+            mixed.float().reshape(B * T, -1),
+            values.float().reshape(B * T, -1),
+            dim=-1,
+        ).mean()
+        return {
+            "token_effect": token_effect,
+            "interaction_effect": interaction_effect,
+            "probe_effect": interaction_effect.reshape(probe_count, -1).mean(dim=-1),
+            "entropy": entropy.mean().item(),
+            "identity_coefficient": identity_coefficient.item(),
+            "dominant_mass": probs.max(dim=-1).values.mean().item(),
+            "head_similarity": head_similarity.item(),
+            "delta_rms": delta_attn.float().square().mean().sqrt().item(),
+            "mix_ratio": (
+                mix_delta.float().norm()
+                / values.float().norm().clamp_min(1e-9)
+            ).item(),
+            "retention_cosine": retention_cosine.item(),
+            "mix_gain": (
+                mixed.float().norm()
+                / values.float().norm().clamp_min(1e-9)
+            ).item(),
+            "key_scale_abs_mean": self.k_scale.float().abs().mean().item(),
+            "logit_max": logits.float().abs().max().item(),
+        }
+
+
+class Block(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        d = config.d_model
+        self.n1, self.n2 = RMSNorm(d), RMSNorm(d)
+        self.attn = TokenAttention(d, config.n_heads, config.n_kv_heads)
+        self.mix = ACSwiGLU(
+            d,
+            config.chunk,
+            config.ac_heads,
+            config.expand,
+            config.ac_mix_scale,
+        )
+
+    def forward(self, x, cos, sin):
+        x = x + self.attn(self.n1(x), cos, sin)
+        x = x + self.mix(self.n2(x))
+        return x
+
+
+class ACSwiGLUModel(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        d = config.d_model
+        self.config = config
+        self.emb = nn.Embedding(config.vocab_size, d)
+        self.blocks = nn.ModuleList(Block(config) for _ in range(config.n_layers))
+        self.norm = RMSNorm(d)
+        hd = d // config.n_heads
+        cos, sin = rope_cache(config.seq_len, hd, "cpu")
+        self.register_buffer("cos", cos)
+        self.register_buffer("sin", sin)
+
+    def forward_hidden(self, idx):
+        if idx.size(1) > self.config.seq_len:
+            idx = idx[:, -self.config.seq_len :]
+        x = self.emb(idx)
+        device_type = x.device.type
+        compute_dtype = (
+            torch.get_autocast_dtype(device_type)
+            if torch.is_autocast_enabled(device_type)
+            else x.dtype
+        )
+        x = x.to(dtype=compute_dtype)
+        cos = self.cos[: idx.size(1)].to(device=idx.device, dtype=compute_dtype)
+        sin = self.sin[: idx.size(1)].to(device=idx.device, dtype=compute_dtype)
+        for b in self.blocks:
+            x = b(x, cos, sin)
+        return self.norm(x)
+
+
+class ACSwiGLUForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = ACSwiGLUConfig
+    base_model_prefix = "model"
+    _no_split_modules = ["Block"]
+    _tied_weights_keys = {"head.weight": "model.emb.weight"}
+    # Transformers 4.x reads the expanded map directly; Transformers 5.x
+    # replaces it during post_init(). Keeping both forms makes the same remote
+    # code load cleanly across that boundary.
+    all_tied_weights_keys = {"head.weight": "model.emb.weight"}
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.model = ACSwiGLUModel(config)
+        self.head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        # Modern Transformers creates loader metadata and performs configured
+        # tying in post_init(); omitting it leaves all_tied_weights_keys absent.
+        self.post_init()
+        self.head.weight = self.model.emb.weight
+
+    def get_input_embeddings(self):
+        return self.model.emb
+
+    def set_input_embeddings(self, value):
+        self.model.emb = value
+
+    def get_output_embeddings(self):
+        return self.head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.head = new_embeddings
+
+    def raw_logits(self, idx):
+        return self.head(self.model.forward_hidden(idx))
+
+    def logits(self, idx):
+        return self.raw_logits(idx)
+
+    def _masked_logits(self, input_ids, attention_mask):
+        if attention_mask is None or bool(attention_mask.all()):
+            return self.logits(input_ids)
+
+        B, T = input_ids.shape
+        out = None
+        for i in range(B):
+            keep = attention_mask[i].bool().nonzero(as_tuple=False).flatten()
+            if keep.numel() == 0:
+                keep = torch.tensor([T - 1], device=input_ids.device)
+            trimmed = input_ids[i, keep].unsqueeze(0)
+            logits_i = self.logits(trimmed)
+            if out is None:
+                out = logits_i.new_zeros(B, T, logits_i.size(-1))
+            out[i, keep, :] = logits_i[0, -keep.numel() :, :]
+        return out
+
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        labels=None,
+        use_cache=False,
+        past_key_values=None,
+        **kwargs,
+    ):
+        return_dict = kwargs.pop(
+            "return_dict", getattr(self.config, "use_return_dict", True)
+        )
+        if input_ids is None:
+            raise ValueError("input_ids must be provided.")
+        if input_ids.size(1) > self.config.seq_len:
+            input_ids = input_ids[:, -self.config.seq_len :]
+            if attention_mask is not None:
+                attention_mask = attention_mask[:, -self.config.seq_len :]
+            if labels is not None:
+                labels = labels[:, -self.config.seq_len :]
+        logits = self._masked_logits(input_ids, attention_mask)
+        loss = None
+        if labels is not None:
+            shift_logits = logits[:, :-1, :].contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)).float(),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+        if not return_dict:
+            return (loss, logits) if loss is not None else (logits,)
+        return CausalLMOutputWithPast(loss=loss, logits=logits, past_key_values=None)
+
+    def prepare_inputs_for_generation(self, input_ids, **kwargs):
+        attention_mask = kwargs.get("attention_mask")
+        result = {
+            "input_ids": input_ids[:, -self.config.seq_len :],
+            "use_cache": False,
+        }
+        if attention_mask is not None:
+            result["attention_mask"] = attention_mask[:, -self.config.seq_len :]
+        return result

tokenizer_config.json

@@ -0,0 +1,14 @@
+{
+  "add_prefix_space": true,
+  "backend": "tokenizers",
+  "bos_token": "<bos>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<eos>",
+  "is_local": false,
+  "local_files_only": false,
+  "model_max_length": 1024,
+  "pad_token": "<eos>",
+  "tokenizer_class": "TokenizersBackend",
+  "unk_token": "<unk>",
+  "vocab_size": 4096
+}

training_config.json

@@ -0,0 +1,64 @@
+{
+  "steps": 40000,
+  "session_steps": 40000,
+  "session_count": 1,
+  "batch_size": 512,
+  "grad_accum": 1,
+  "seq_len": 1024,
+  "model_context": 1024,
+  "d_model": 216,
+  "n_layers": 10,
+  "n_heads": 6,
+  "n_kv_heads": 2,
+  "chunk": 24,
+  "ac_heads": 3,
+  "expand": 2,
+  "ac_mix_scale": 1.0,
+  "lr": 0.0025,
+  "linear_decay_start": 30000,
+  "cosine_decay_start": 30000,
+  "mid_lr": 0.0025,
+  "muon_lr": 0.03,
+  "muon_momentum": 0.95,
+  "muon_ns_steps": 5,
+  "muon_adjust_lr_fn": null,
+  "warmup": 1000,
+  "wd": 0.01,
+  "grad_clip": 1.0,
+  "log_every": 10,
+  "diag_every": 1000,
+  "eval_every": 10000,
+  "val_batch_size": 32,
+  "val_batches": 0,
+  "val_context": 1024,
+  "val_stride": 512,
+  "infer_tokens": 512,
+  "infer_repeat_penalty": 1.2,
+  "infer_prompt": "The process of photosynthesis",
+  "lm_eval_tasks": "arc_easy,arc_challenge,hellaswag,piqa",
+  "lm_eval_batch_size": "32",
+  "lm_eval_device": "cuda",
+  "lm_eval_dtype": "bfloat16",
+  "lm_eval_softmax_dtype": "float32",
+  "lm_eval_expected_version": "0.4.12",
+  "lm_eval_retries": 3,
+  "lm_eval_export_dir": "AC_SwiGLU_Inclusive_40k_lm_eval_hf",
+  "lm_eval_output_dir": "lm_eval_results_AC_SwiGLU_Inclusive_40k",
+  "arithmark3_choice_batch_size": 64,
+  "arithmark3_force_download": false,
+  "recipe_version": "ACSwiGLUInclusive_d216_l10_q6_kv2_chunk24_h3_expand2_centered18x18_mix1_train1024_ctx1024_v4k_gpts4944k_fullcompile_adam2p5e3_hold30k_cos0_40k_muon3e2_w1k_clip1_20p97B_b512_a1_tpu524k_session1x40k_static_fwe55_cos25_fwhq10_math10_interleaved_pinneddata_buf1k_prefetch16_eval1024s512_b32_arithb64_eval10k_diag1k_arithbos_scorew75_v1",
+  "hf_repo_id": "User01110/attention-contraction-mini",
+  "hf_repo_private": false,
+  "hub_upload_retries": 3,
+  "resume_branch": "resume-latest",
+  "resume_model_file": "resume_model.safetensors",
+  "resume_state_file": "resume_state.pt",
+  "resume_manifest_file": "resume_manifest.json",
+  "tokenizer_name": "AxiomicLabs/GPT-S-5M",
+  "tokenizer_revision": "275b9c3ca78736bf6aeb154c7e2d5f5764fe9035",
+  "data_seed": 1337,
+  "shuffle_buffer": 1024,
+  "tokenize_batch_size": 64,
+  "prefetch_batches": 16,
+  "compile": true
+}