Publish Tessera-1B-Nano checkpoint

README.md

@@ -0,0 +1,86 @@
+---
+library_name: transformers
+pipeline_tag: text-generation
+base_model: HuggingFaceTB/SmolLM2-135M
+tags:
+- next-concept-prediction
+- conceptlm
+- causal-lm
+- smollm2
+- tessera
+license: apache-2.0
+---
+
+# Tessera-1B-Nano
+
+A Next Concept Prediction checkpoint from Paragon Intelligence Labs: an independent,
+small-scale replication of ConceptLM on `HuggingFaceTB/SmolLM2-135M`. The model keeps
+ordinary next-token generation and adds a causal, product-quantized concept path over
+4-token chunks: token states are pooled into concepts,
+product-quantized, processed by causal concept blocks, and the predicted next concept
+is fed back into the token decoder.
+
+## TL;DR
+
+This checkpoint comes from a token-matched comparison against the unchanged
+HuggingFaceTB/SmolLM2-135M backbone: both arms consumed the same tokens of the
+same packed corpus, in the same order, from the same initialization, with no
+restarts and no NaNs.
+
+- **Token loss is neutral**: the matched baseline is within run noise (see the
+  whitepaper and `results/fineweb-edu/README.md` in the ncp-smol repository for
+  the exact numbers).
+- **The concept channel is causally used**: zeroing predicted concept feedback
+  at inference costs +0.0205 nats of held-out NTP loss.
+- **The codebook is rich**: no low-entropy shortcut, usage grows monotonically
+  over the run (growth curves are in the repository whitepaper).
+- Full reading of the interventions: `docs/whitepaper.md` in the ncp-smol repo.
+
+## Evaluation
+
+| Metric | Value |
+| --- | ---: |
+| Held-out NTP loss | 1.8751 |
+| Held-out perplexity | 6.5213 |
+| Zero feedback delta | 0.0205 |
+| Shuffled feedback delta | 0.0000 |
+| Codebook perplexity / usage | 2.4087 / 33.0% |
+| Training tokens | 16,777,216 |
+| Tracked compute estimate | $1.38 |
+
+Intervention deltas are increases in held-out NTP loss relative to normal predicted concept
+feedback, evaluated on identical batches.
+
+## Architecture
+
+- chunk size: 4
+- product code: 9 segments x 64 entries
+- causal concept blocks: 2
+- injection point: before token decoder block 2
+- NCP target: next continuous concept
+- loss: `L_ntp + 1 L_ncp + 1 L_vq`
+
+This is a compact ConceptLM-style implementation, not an 8.9B NCP-ArchPreview replica. It
+omits iterative residual coding, cross-scale residual connections, and the large-scale
+training recipe.
+
+## Training data and provenance
+
+The matched corpus and the training record are pinned in the ncp-smol repository:
+packed-cache SHA256 hashes, the complete metric log, trainer state, and the raw
+intervention evaluation JSON (also shipped in this repository as `eval.json` and
+`metrics.jsonl`).
+
+## Loading
+
+```python
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+tokenizer = AutoTokenizer.from_pretrained("yava-code/Tessera-135M-Gate")
+model = AutoModelForCausalLM.from_pretrained("yava-code/Tessera-135M-Gate")
+```
+
+## References
+
+- ConceptLM: https://arxiv.org/abs/2602.08984
+- NCP-ArchPreview: https://arxiv.org/abs/2609.10715

config.json

@@ -0,0 +1,115 @@
+{
+  "architectures": [
+    "NcpSmolForCausalLM"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration.NcpSmolConfig",
+    "AutoModelForCausalLM": "modeling.NcpSmolForCausalLM"
+  },
+  "backbone_config": {
+    "_name_or_path": "HuggingFaceTB/SmolLM2-135M",
+    "add_cross_attention": false,
+    "architectures": [
+      "LlamaForCausalLM"
+    ],
+    "attention_bias": false,
+    "attention_dropout": 0.0,
+    "attn_implementation": null,
+    "bad_words_ids": null,
+    "begin_suppress_tokens": null,
+    "bos_token_id": 0,
+    "chunk_size_feed_forward": 0,
+    "cross_attention_hidden_size": null,
+    "decoder_start_token_id": null,
+    "diversity_penalty": 0.0,
+    "do_sample": false,
+    "dtype": "float32",
+    "early_stopping": false,
+    "encoder_no_repeat_ngram_size": 0,
+    "eos_token_id": 0,
+    "exponential_decay_length_penalty": null,
+    "finetuning_task": null,
+    "forced_bos_token_id": null,
+    "forced_eos_token_id": null,
+    "head_dim": 64,
+    "hidden_act": "silu",
+    "hidden_size": 576,
+    "id2label": {
+      "0": "LABEL_0",
+      "1": "LABEL_1"
+    },
+    "initializer_range": 0.041666666666666664,
+    "intermediate_size": 1536,
+    "is_decoder": false,
+    "is_encoder_decoder": false,
+    "is_llama_config": true,
+    "label2id": {
+      "LABEL_0": 0,
+      "LABEL_1": 1
+    },
+    "length_penalty": 1.0,
+    "max_length": 20,
+    "max_position_embeddings": 8192,
+    "min_length": 0,
+    "mlp_bias": false,
+    "model_type": "llama",
+    "no_repeat_ngram_size": 0,
+    "num_attention_heads": 9,
+    "num_beam_groups": 1,
+    "num_beams": 1,
+    "num_hidden_layers": 30,
+    "num_key_value_heads": 3,
+    "num_return_sequences": 1,
+    "output_attentions": false,
+    "output_hidden_states": false,
+    "output_scores": false,
+    "pad_token_id": null,
+    "prefix": null,
+    "pretraining_tp": 1,
+    "problem_type": null,
+    "pruned_heads": {},
+    "remove_invalid_values": false,
+    "repetition_penalty": 1.0,
+    "return_dict": true,
+    "return_dict_in_generate": false,
+    "rms_norm_eps": 1e-05,
+    "rope_interleaved": false,
+    "rope_scaling": null,
+    "rope_theta": 100000,
+    "sep_token_id": null,
+    "suppress_tokens": null,
+    "task_specific_params": null,
+    "temperature": 1.0,
+    "tf_legacy_loss": false,
+    "tie_encoder_decoder": false,
+    "tie_word_embeddings": true,
+    "tokenizer_class": null,
+    "top_k": 50,
+    "top_p": 1.0,
+    "torchscript": false,
+    "transformers_version": "4.57.6",
+    "typical_p": 1.0,
+    "use_bfloat16": false,
+    "use_cache": true,
+    "vocab_size": 49152
+  },
+  "base_model_name_or_path": "HuggingFaceTB/SmolLM2-135M",
+  "base_model_revision": "93efa2f097d58c2a74874c7e644dbc9b0cee75a2",
+  "chunk_size": 4,
+  "codebook_size": 64,
+  "concept_layers": 2,
+  "dtype": "float32",
+  "hidden_size": 576,
+  "insert_layer": 1,
+  "model_type": "ncp_smol",
+  "ncp_target": "continuous",
+  "ncp_weight": 1.0,
+  "num_attention_heads": 9,
+  "num_hidden_layers": 30,
+  "num_key_value_heads": 3,
+  "segments": 9,
+  "transformers_version": "4.57.6",
+  "use_cache": false,
+  "vocab_size": 49152,
+  "vq_weight": 1.0
+}

configuration.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+from typing import Any
+
+from transformers import LlamaConfig, PretrainedConfig
+
+
+class NcpSmolConfig(PretrainedConfig):
+    model_type = "ncp_smol"
+
+    def __init__(
+        self,
+        backbone_config: dict[str, Any] | None = None,
+        base_model_name_or_path: str = "HuggingFaceTB/SmolLM2-360M",
+        base_model_revision: str | None = None,
+        chunk_size: int = 4,
+        segments: int | None = None,
+        codebook_size: int = 64,
+        concept_layers: int = 2,
+        insert_layer: int = 1,
+        ncp_target: str = "continuous",
+        ncp_weight: float = 1.0,
+        vq_weight: float = 1.0,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        backbone_config = backbone_config or LlamaConfig().to_dict()
+        hidden_size = int(backbone_config["hidden_size"])
+        heads = int(backbone_config["num_attention_heads"])
+        segments = heads if segments is None else segments
+
+        if chunk_size < 2:
+            raise ValueError("chunk_size must be at least 2")
+        if hidden_size % segments:
+            raise ValueError("hidden_size must be divisible by segments")
+        if not 0 < insert_layer < int(backbone_config["num_hidden_layers"]):
+            raise ValueError("insert_layer must split the token backbone")
+        if concept_layers < 1:
+            raise ValueError("concept_layers must be positive")
+        if codebook_size < 2:
+            raise ValueError("codebook_size must be at least 2")
+        if ncp_target not in {"quantized", "continuous"}:
+            raise ValueError("ncp_target must be 'quantized' or 'continuous'")
+
+        self.backbone_config = backbone_config
+        self.hidden_size = hidden_size
+        self.vocab_size = int(backbone_config["vocab_size"])
+        self.num_hidden_layers = int(backbone_config["num_hidden_layers"])
+        self.num_attention_heads = heads
+        self.num_key_value_heads = int(backbone_config["num_key_value_heads"])
+        self.use_cache = False
+        self.base_model_name_or_path = base_model_name_or_path
+        self.base_model_revision = base_model_revision
+        self.chunk_size = chunk_size
+        self.segments = segments
+        self.codebook_size = codebook_size
+        self.concept_layers = concept_layers
+        self.insert_layer = insert_layer
+        self.ncp_target = ncp_target
+        self.ncp_weight = ncp_weight
+        self.vq_weight = vq_weight
+        self.architectures = ["NcpSmolForCausalLM"]
+        self.auto_map = {
+            "AutoConfig": "configuration.NcpSmolConfig",
+            "AutoModelForCausalLM": "modeling.NcpSmolForCausalLM",
+        }
+
+
+NcpSmolConfig.register_for_auto_class()

data_metadata.json

@@ -0,0 +1,14 @@
+{
+  "dataset": "roneneldan/TinyStories",
+  "dtype": "uint32",
+  "eos_token_id": 0,
+  "revision": "f54c09fd23315a6f9c86f9dc80f725de7d8f9c64",
+  "sequence_length": 256,
+  "subset": null,
+  "tokenizer": "HuggingFaceTB/SmolLM2-135M",
+  "tokenizer_revision": "93efa2f097d58c2a74874c7e644dbc9b0cee75a2",
+  "train_sha256": "1b72ac9f0a05d841df30feb6feb8fbd7c38170d61528e78cfea3b04c491cc025",
+  "train_tokens": 1048576,
+  "validation_sha256": "6f0c6b7c842665a23d593f34c00ac7d7a9a9b3050a3d1d1592fac27b12be2462",
+  "validation_tokens": 1048576
+}

eval.json

@@ -0,0 +1,55 @@
+{
+  "intervention": {
+    "shuffle_minus_predicted": 2.2351741790771484e-08,
+    "similar_shuffle_minus_predicted": -2.9802322387695312e-08,
+    "zero_minus_predicted": 0.020453326404094696
+  },
+  "predicted": {
+    "codebook_perplexity": 2.4087092876434326,
+    "codebook_usage": 0.3298611044883728,
+    "loss_by_chunk_offset": {
+      "0": 1.9171534776687622,
+      "1": 1.861592411994934,
+      "2": 1.8690117597579956,
+      "3": 1.8531707525253296
+    },
+    "ntp_loss": 1.8750678300857544,
+    "perplexity": 6.521261443051458
+  },
+  "shuffle": {
+    "codebook_perplexity": 2.4087092876434326,
+    "codebook_usage": 0.3298611044883728,
+    "loss_by_chunk_offset": {
+      "0": 1.9171533584594727,
+      "1": 1.8615922927856445,
+      "2": 1.869011640548706,
+      "3": 1.8531708717346191
+    },
+    "ntp_loss": 1.8750678524374962,
+    "perplexity": 6.521261588813012
+  },
+  "similar_shuffle": {
+    "codebook_perplexity": 2.4087092876434326,
+    "codebook_usage": 0.3298611044883728,
+    "loss_by_chunk_offset": {
+      "0": 1.9171534776687622,
+      "1": 1.861592411994934,
+      "2": 1.8690117597579956,
+      "3": 1.8531707525253296
+    },
+    "ntp_loss": 1.875067800283432,
+    "perplexity": 6.521261248702725
+  },
+  "zero": {
+    "codebook_perplexity": 2.4087092876434326,
+    "codebook_usage": 0.3298611044883728,
+    "loss_by_chunk_offset": {
+      "0": 1.9391287565231323,
+      "1": 1.896687388420105,
+      "2": 1.8889083862304688,
+      "3": 1.8580405712127686
+    },
+    "ntp_loss": 1.895521156489849,
+    "perplexity": 6.656016326994612
+  }
+}

experiment.json

@@ -0,0 +1,60 @@
+{
+  "data": {
+    "cache_dir": "/vol/.cache/data/tinystories-overfit",
+    "cache_train_tokens": 1048576,
+    "dataset": "roneneldan/TinyStories",
+    "overfit": true,
+    "revision": "f54c09fd23315a6f9c86f9dc80f725de7d8f9c64",
+    "sequence_length": 256,
+    "shuffle_buffer": 2000,
+    "subset": null,
+    "text_column": "text",
+    "train_split": "train",
+    "train_tokens": 16777216,
+    "validation_split": "validation",
+    "validation_tokens": 1048576
+  },
+  "model": {
+    "base_model": "HuggingFaceTB/SmolLM2-135M",
+    "chunk_size": 4,
+    "codebook_size": 64,
+    "concept_layers": 2,
+    "dtype": "bfloat16",
+    "insert_layer": 1,
+    "ncp_target": "continuous",
+    "ncp_weight": 1.0,
+    "revision": "93efa2f097d58c2a74874c7e644dbc9b0cee75a2",
+    "segments": 9,
+    "vq_weight": 1.0
+  },
+  "optim": {
+    "beta1": 0.9,
+    "beta2": 0.95,
+    "eps": 1e-08,
+    "grad_accum_steps": 4,
+    "learning_rate": 3e-05,
+    "max_grad_norm": 1.0,
+    "micro_batch_size": 4,
+    "min_lr_ratio": 0.1,
+    "warmup_ratio": 0.01,
+    "weight_decay": 0.1
+  },
+  "run": {
+    "mode": "ncp",
+    "name": "tinystories-overfit",
+    "output_dir": "/vol/runs/tinystories-overfit",
+    "resume": true,
+    "seed": 17
+  },
+  "train": {
+    "compile": false,
+    "cost_overhead": 1.15,
+    "eval_batches": 16,
+    "eval_every_tokens": 4194304,
+    "gpu_hourly_usd": 0.8,
+    "log_every_steps": 10,
+    "max_wall_time_minutes": 180,
+    "run_budget_usd": 6.0,
+    "save_every_tokens": 8388608
+  }
+}

generation_config.json

@@ -0,0 +1,5 @@
+{
+  "_from_model_config": true,
+  "transformers_version": "4.57.6",
+  "use_cache": false
+}

model.safetensors

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

modeling.py

@@ -0,0 +1,334 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Literal
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor, nn
+from transformers import AutoModelForCausalLM, GenerationMixin, LlamaConfig, PreTrainedModel
+from transformers.modeling_outputs import ModelOutput
+from transformers.models.llama.modeling_llama import (
+    LlamaDecoderLayer,
+    LlamaForCausalLM,
+    LlamaRotaryEmbedding,
+)
+
+from .configuration import NcpSmolConfig
+from .quantizer import ProductVectorQuantizer
+
+ConceptMode = Literal["predicted", "zero", "shuffle"]
+
+
+@dataclass
+class NcpCausalLMOutput(ModelOutput):
+    loss: Tensor | None = None
+    logits: Tensor | None = None
+    ntp_loss: Tensor | None = None
+    ncp_loss: Tensor | None = None
+    vq_loss: Tensor | None = None
+    code_indices: Tensor | None = None
+
+
+class NcpSmolForCausalLM(PreTrainedModel, GenerationMixin):
+    config_class = NcpSmolConfig
+    base_model_prefix = "backbone"
+    main_input_name = "input_ids"
+    _supports_cache_class = False
+    _no_split_modules = ["LlamaDecoderLayer"]
+    _tied_weights_keys = ["backbone.lm_head.weight"]
+
+    def __init__(self, config: NcpSmolConfig) -> None:
+        super().__init__(config)
+        backbone_config = LlamaConfig.from_dict(config.backbone_config)
+        self.backbone = LlamaForCausalLM(backbone_config)
+
+        concept_config = LlamaConfig.from_dict(config.backbone_config)
+        concept_config._attn_implementation = "eager"
+        self.concept_layers = nn.ModuleList(
+            [LlamaDecoderLayer(concept_config, i) for i in range(config.concept_layers)]
+        )
+        self.concept_rotary = LlamaRotaryEmbedding(concept_config)
+        self.quantizer = ProductVectorQuantizer(
+            hidden_size=config.hidden_size,
+            segments=config.segments,
+            codebook_size=config.codebook_size,
+        )
+        self.concept_head = nn.Linear(
+            config.hidden_size,
+            config.segments * config.codebook_size,
+        )
+        self._hook_state: dict[str, Any] = {}
+        self._hook = self.backbone.model.layers[config.insert_layer].register_forward_pre_hook(
+            self._inject_concepts,
+            with_kwargs=True,
+        )
+        self.post_init()
+        self.init_concept_from_backbone()
+        self.tie_weights()
+
+    @classmethod
+    def from_backbone(
+        cls,
+        model_name_or_path: str,
+        *,
+        revision: str | None = None,
+        chunk_size: int = 4,
+        segments: int | None = None,
+        codebook_size: int = 64,
+        concept_layers: int = 2,
+        insert_layer: int = 1,
+        ncp_target: str = "continuous",
+        ncp_weight: float = 1.0,
+        vq_weight: float = 1.0,
+        **load_kwargs: Any,
+    ) -> NcpSmolForCausalLM:
+        base = AutoModelForCausalLM.from_pretrained(
+            model_name_or_path,
+            revision=revision,
+            **load_kwargs,
+        )
+        if not isinstance(base, LlamaForCausalLM):
+            raise TypeError("ncp-smol currently requires a Llama-compatible backbone")
+        config = NcpSmolConfig(
+            backbone_config=base.config.to_dict(),
+            base_model_name_or_path=model_name_or_path,
+            base_model_revision=revision,
+            chunk_size=chunk_size,
+            segments=segments,
+            codebook_size=codebook_size,
+            concept_layers=concept_layers,
+            insert_layer=insert_layer,
+            ncp_target=ncp_target,
+            ncp_weight=ncp_weight,
+            vq_weight=vq_weight,
+        )
+        model = cls(config)
+        model.backbone.load_state_dict(base.state_dict())
+        model.init_concept_from_backbone()
+        model.tie_weights()
+        return model
+
+    def init_concept_from_backbone(self) -> None:
+        source = self.backbone.model.layers
+        for index, layer in enumerate(self.concept_layers):
+            layer.load_state_dict(source[index % len(source)].state_dict())
+
+    def get_input_embeddings(self) -> nn.Module:
+        return self.backbone.get_input_embeddings()
+
+    def set_input_embeddings(self, value: nn.Module) -> None:
+        self.backbone.set_input_embeddings(value)
+
+    def get_output_embeddings(self) -> nn.Module:
+        return self.backbone.get_output_embeddings()
+
+    def set_output_embeddings(self, value: nn.Module) -> None:
+        self.backbone.set_output_embeddings(value)
+
+    def tie_weights(self) -> None:
+        self.backbone.tie_weights()
+
+    def _concept_mask(self, valid: Tensor, dtype: torch.dtype) -> Tensor:
+        batch, length = valid.shape
+        mask = torch.full(
+            (length, length),
+            torch.finfo(dtype).min,
+            dtype=dtype,
+            device=valid.device,
+        ).triu(diagonal=1)
+        mask = mask.view(1, 1, length, length).expand(batch, 1, length, length).clone()
+        return mask.masked_fill(~valid[:, None, None, :], torch.finfo(dtype).min)
+
+    def _concept_path(
+        self,
+        hidden: Tensor,
+        attention_mask: Tensor | None,
+        mode: ConceptMode,
+    ) -> tuple[Tensor, dict[str, Tensor]]:
+        batch, length, width = hidden.shape
+        chunks = length // self.config.chunk_size
+        zero = hidden.sum() * 0
+        if chunks == 0:
+            return torch.zeros_like(hidden), {
+                "ncp_loss": zero,
+                "vq_loss": zero,
+                "code_indices": torch.empty(
+                    batch, 0, self.config.segments, dtype=torch.long, device=hidden.device
+                ),
+            }
+
+        usable = chunks * self.config.chunk_size
+        concepts = (
+            hidden[:, :usable]
+            .reshape(
+                batch,
+                chunks,
+                self.config.chunk_size,
+                width,
+            )
+            .mean(dim=2)
+        )
+        if attention_mask is None:
+            valid = torch.ones(batch, chunks, dtype=torch.bool, device=hidden.device)
+        else:
+            valid = (
+                attention_mask[:, :usable]
+                .bool()
+                .reshape(batch, chunks, self.config.chunk_size)
+                .all(dim=2)
+            )
+
+        quantized = self.quantizer(concepts)
+        positions = torch.arange(chunks, device=hidden.device)
+        position_ids = (positions * self.config.chunk_size).unsqueeze(0).expand(batch, chunks)
+        concept_hidden = concepts
+        rotary = self.concept_rotary(concept_hidden, position_ids)
+        causal_mask = self._concept_mask(valid, concept_hidden.dtype)
+        for layer in self.concept_layers:
+            concept_hidden = layer(
+                concept_hidden,
+                attention_mask=causal_mask,
+                position_ids=position_ids,
+                cache_position=positions,
+                position_embeddings=rotary,
+                use_cache=False,
+            )
+
+        logits = self.concept_head(concept_hidden).view(
+            batch,
+            chunks,
+            self.config.segments,
+            self.config.codebook_size,
+        )
+        predicted = self.quantizer.expected(logits)
+
+        if chunks > 1 and valid[:, 1:].any():
+            target_mask = valid[:, 1:].unsqueeze(-1)
+            target = quantized.codes if self.config.ncp_target == "quantized" else concepts
+            squared = (predicted[:, :-1] - target[:, 1:].detach()).square()
+            ncp_loss = (squared * target_mask).sum() / (target_mask.sum().clamp_min(1) * width)
+        else:
+            ncp_loss = zero
+
+        if mode == "zero":
+            predicted = torch.zeros_like(predicted)
+        elif mode == "shuffle":
+            predicted = predicted.roll(shifts=1, dims=0)
+
+        feedback = self._align_feedback(predicted.detach(), length)
+        return feedback, {
+            "ncp_loss": ncp_loss,
+            "vq_loss": quantized.loss,
+            "code_indices": quantized.indices,
+        }
+
+    def _align_feedback(self, predicted: Tensor, length: int) -> Tensor:
+        batch, _, width = predicted.shape
+        prefix = torch.zeros(batch, 1, width, dtype=predicted.dtype, device=predicted.device)
+        repeated = torch.cat((prefix, predicted), dim=1).repeat_interleave(
+            self.config.chunk_size,
+            dim=1,
+        )
+        return repeated[:, 1 : length + 1]
+
+    def _inject_concepts(
+        self,
+        _module: nn.Module,
+        args: tuple[Any, ...],
+        kwargs: dict[str, Any],
+    ) -> tuple[tuple[Any, ...], dict[str, Any]]:
+        if not self._hook_state.get("active", False):
+            return args, kwargs
+
+        hidden = args[0] if args else kwargs["hidden_states"]
+        feedback, aux = self._concept_path(
+            hidden,
+            self._hook_state.get("attention_mask"),
+            self._hook_state["mode"],
+        )
+        self._hook_state["aux"] = aux
+        if args:
+            args = (hidden + feedback, *args[1:])
+        else:
+            kwargs["hidden_states"] = hidden + feedback
+        return args, kwargs
+
+    def forward(
+        self,
+        input_ids: Tensor | None = None,
+        attention_mask: Tensor | None = None,
+        labels: Tensor | None = None,
+        inputs_embeds: Tensor | None = None,
+        concept_mode: ConceptMode = "predicted",
+        use_cache: bool | None = None,
+        return_dict: bool = True,
+        **kwargs: Any,
+    ) -> NcpCausalLMOutput | tuple[Tensor, ...]:
+        if concept_mode not in {"predicted", "zero", "shuffle"}:
+            raise ValueError(f"unknown concept_mode: {concept_mode}")
+
+        self._hook_state = {
+            "active": True,
+            "attention_mask": attention_mask,
+            "mode": concept_mode,
+            "aux": None,
+        }
+        try:
+            outputs = self.backbone(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                inputs_embeds=inputs_embeds,
+                labels=None,
+                use_cache=False,
+                return_dict=True,
+                **kwargs,
+            )
+            aux = self._hook_state["aux"]
+        finally:
+            self._hook_state["active"] = False
+        if aux is None:
+            raise RuntimeError("concept injection hook did not run")
+
+        ntp_loss = None
+        total_loss = None
+        if labels is not None:
+            shift_logits = outputs.logits[:, :-1].float().contiguous()
+            shift_labels = labels[:, 1:].contiguous()
+            ntp_loss = F.cross_entropy(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1),
+                ignore_index=-100,
+            )
+            total_loss = (
+                ntp_loss
+                + self.config.ncp_weight * aux["ncp_loss"]
+                + self.config.vq_weight * aux["vq_loss"]
+            )
+
+        result = NcpCausalLMOutput(
+            loss=total_loss,
+            logits=outputs.logits,
+            ntp_loss=ntp_loss,
+            ncp_loss=aux["ncp_loss"],
+            vq_loss=aux["vq_loss"],
+            code_indices=aux["code_indices"],
+        )
+        if return_dict:
+            return result
+        return tuple(value for value in result.values() if value is not None)
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: Tensor,
+        attention_mask: Tensor | None = None,
+        **_: Any,
+    ) -> dict[str, Any]:
+        return {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "use_cache": False,
+        }
+
+
+NcpSmolForCausalLM.register_for_auto_class("AutoModelForCausalLM")

quantizer.py

@@ -0,0 +1,100 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor, nn
+
+
+@dataclass
+class QuantizerOutput:
+    codes: Tensor
+    indices: Tensor
+    loss: Tensor
+
+
+class CodebookTransform(nn.Module):
+    def __init__(self, segments: int, dim: int) -> None:
+        super().__init__()
+        self.layers = nn.ModuleList(
+            [
+                nn.Sequential(
+                    nn.Linear(dim, 2 * dim),
+                    nn.ReLU(),
+                    nn.Linear(2 * dim, dim),
+                )
+                for _ in range(segments)
+            ]
+        )
+
+    def forward(self, codebook: Tensor) -> Tensor:
+        return torch.stack(
+            [layer(codes) for layer, codes in zip(self.layers, codebook, strict=True)]
+        )
+
+
+class ProductVectorQuantizer(nn.Module):
+    def __init__(self, hidden_size: int, segments: int, codebook_size: int) -> None:
+        super().__init__()
+        if hidden_size % segments:
+            raise ValueError("hidden_size must be divisible by segments")
+
+        self.hidden_size = hidden_size
+        self.segments = segments
+        self.codebook_size = codebook_size
+        self.segment_dim = hidden_size // segments
+        self.register_buffer(
+            "codebook",
+            torch.empty(segments, codebook_size, self.segment_dim),
+        )
+        self.transform = CodebookTransform(segments, self.segment_dim)
+        nn.init.normal_(self.codebook, std=0.02)
+
+    def transformed_codes(self) -> Tensor:
+        return self.transform(self.codebook)
+
+    def forward(self, concepts: Tensor) -> QuantizerOutput:
+        if concepts.shape[-1] != self.hidden_size:
+            raise ValueError(f"expected hidden size {self.hidden_size}, got {concepts.shape[-1]}")
+
+        shape = concepts.shape[:-1]
+        x = concepts.reshape(*shape, self.segments, self.segment_dim)
+        target = x.detach()
+        codes = self.transformed_codes()
+        distances = (
+            target.square().sum(dim=-1, keepdim=True)
+            + codes.square()
+            .sum(dim=-1)
+            .view(*([1] * len(shape)), self.segments, self.codebook_size)
+            - 2 * torch.einsum("...sd,snd->...sn", target, codes)
+        )
+        indices = distances.argmin(dim=-1)
+        gather_index = indices.unsqueeze(-1).expand(*indices.shape, self.segment_dim)
+        expanded = codes.view(*([1] * len(shape)), *codes.shape).expand(*shape, *codes.shape)
+        quantized = expanded.gather(-2, gather_index.unsqueeze(-2)).squeeze(-2)
+        loss = F.mse_loss(quantized, target)
+        return QuantizerOutput(quantized.reshape(*shape, self.hidden_size), indices, loss)
+
+    def expected(self, logits: Tensor) -> Tensor:
+        if logits.shape[-2:] != (self.segments, self.codebook_size):
+            raise ValueError("logits do not match the product codebook")
+        probs = logits.float().softmax(dim=-1).to(logits.dtype)
+        predicted = torch.einsum("...sn,snd->...sd", probs, self.transformed_codes())
+        return predicted.flatten(-2)
+
+    @torch.no_grad()
+    def usage(self, indices: Tensor) -> dict[str, float]:
+        flat = indices.reshape(-1, self.segments)
+        perplexities = []
+        active = []
+        for segment in range(self.segments):
+            counts = torch.bincount(flat[:, segment], minlength=self.codebook_size).float()
+            probs = counts / counts.sum().clamp_min(1)
+            entropy = -(probs * probs.clamp_min(1e-12).log()).sum()
+            perplexities.append(entropy.exp())
+            active.append((counts > 0).float().mean())
+        return {
+            "codebook_perplexity": torch.stack(perplexities).mean().item(),
+            "codebook_usage": torch.stack(active).mean().item(),
+        }

special_tokens_map.json

@@ -0,0 +1,42 @@
+{
+  "additional_special_tokens": [
+    "<|endoftext|>",
+    "<|im_start|>",
+    "<|im_end|>",
+    "<repo_name>",
+    "<reponame>",
+    "<file_sep>",
+    "<filename>",
+    "<gh_stars>",
+    "<issue_start>",
+    "<issue_comment>",
+    "<issue_closed>",
+    "<jupyter_start>",
+    "<jupyter_text>",
+    "<jupyter_code>",
+    "<jupyter_output>",
+    "<jupyter_script>",
+    "<empty_output>"
+  ],
+  "bos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<|endoftext|>",
+    "lstrip": false,
+    "normalized": false,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer_config.json

@@ -0,0 +1,168 @@
+{
+  "add_prefix_space": false,
+  "added_tokens_decoder": {
+    "0": {
+      "content": "<|endoftext|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "1": {
+      "content": "<|im_start|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "2": {
+      "content": "<|im_end|>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "3": {
+      "content": "<repo_name>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "4": {
+      "content": "<reponame>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "5": {
+      "content": "<file_sep>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "6": {
+      "content": "<filename>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "7": {
+      "content": "<gh_stars>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "8": {
+      "content": "<issue_start>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "9": {
+      "content": "<issue_comment>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "10": {
+      "content": "<issue_closed>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "11": {
+      "content": "<jupyter_start>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "12": {
+      "content": "<jupyter_text>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "13": {
+      "content": "<jupyter_code>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "14": {
+      "content": "<jupyter_output>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "15": {
+      "content": "<jupyter_script>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    },
+    "16": {
+      "content": "<empty_output>",
+      "lstrip": false,
+      "normalized": false,
+      "rstrip": false,
+      "single_word": false,
+      "special": true
+    }
+  },
+  "additional_special_tokens": [
+    "<|endoftext|>",
+    "<|im_start|>",
+    "<|im_end|>",
+    "<repo_name>",
+    "<reponame>",
+    "<file_sep>",
+    "<filename>",
+    "<gh_stars>",
+    "<issue_start>",
+    "<issue_comment>",
+    "<issue_closed>",
+    "<jupyter_start>",
+    "<jupyter_text>",
+    "<jupyter_code>",
+    "<jupyter_output>",
+    "<jupyter_script>",
+    "<empty_output>"
+  ],
+  "bos_token": "<|endoftext|>",
+  "clean_up_tokenization_spaces": false,
+  "eos_token": "<|endoftext|>",
+  "extra_special_tokens": {},
+  "model_max_length": 8192,
+  "tokenizer_class": "GPT2Tokenizer",
+  "unk_token": "<|endoftext|>",
+  "vocab_size": 49152
+}

trainer_state.json

@@ -0,0 +1,9 @@
+{
+  "batcher": {
+    "cursor": 4096,
+    "epoch": 15
+  },
+  "billable_seconds": 3455.0923419839996,
+  "step": 4096,
+  "tokens_seen": 16777216
+}