--- license: apache-2.0 base_model: HuggingFaceTB/SmolLM2-135M library_name: onnx pipeline_tag: text-generation tags: - transformers.js - onnx - interpretability - attention - hidden-states - tuned-lens --- # SmolLM2-135M Observable This is a weight-only Q4 ONNX export of [`HuggingFaceTB/SmolLM2-135M`](https://huggingface.co/HuggingFaceTB/SmolLM2-135M) for an interactive course laboratory. It exposes the intermediate tensors needed to inspect one forward pass in a browser. The model is a base next-token predictor. It is not an instruction-following or chat model. ## Inputs The sequence inputs use `int64` tensors with shape `[batch, sequence]`: - `input_ids` - `attention_mask` `query_index` is an `int64` tensor with shape `[batch]`. It selects the token position whose layer update will be decomposed. The browser laboratory uses a batch size of one and short sequences. ## Outputs The graph returns 182 tensors: - `next_token_logits`, with shape `[batch, 49152]`; - `hidden_state_00` through `hidden_state_30`, with shape `[batch, sequence, 576]`; - `attention_01` through `attention_30`, with shape `[batch, 9, query sequence, key sequence]`; - `attention_contribution_01` through `attention_contribution_30`, with shape `[batch, sequence, 576]`; - `counterfactual_hidden_state_01` through `counterfactual_hidden_state_30`, with shape `[batch, sequence, 576]`; - `attention_output_01` through `attention_output_30`, with shape `[batch, 576]`; - `mlp_output_01` through `mlp_output_30`, with shape `[batch, 576]`. `hidden_state_00` contains the input embeddings. The following hidden states contain the raw residual stream after each transformer layer, before the final RMSNorm. For a query token `i` and source token `j`, each contribution is: ```text cᵢ⟵ⱼ = Wₒ concatₕ(aᵢⱼʰ vⱼʰ) ``` The source contributions sum to `attention_output`. Adding that result to the previous residual state gives the post-attention state. Adding `mlp_output` gives the next residual state. The manifest records numerical checks for both identities. Each `counterfactual_hidden_state` removes one source contribution before the layer MLP and evaluates that MLP again. Subtracting this state from the complete state gives the source token's effect after the nonlinearity. ## Quantization The published graph uses 4-bit weight-only quantization with asymmetric 32-value blocks. Activations and outputs remain float32. The graph uses the ONNX Runtime `MatMulNBits` operator. The export script compares the ONNX result with the original PyTorch model. The included manifest records the validation prompt, tensor shapes and numerical comparison for the published file. ## Intended use This artifact supports a browser laboratory that visualizes: - next-token probabilities; - token boundaries; - mean attention across heads in the selected layer; - the most likely next token if the residual stream were read at the selected layer. The repository includes two lens graphs. `model_logit_lens_q8.onnx` applies the model's final RMSNorm and original `lm_head` directly to an intermediate residual-stream state. `model_tuned_lens_q8.onnx` first applies an affine translator trained for that layer, then uses the same RMSNorm and `lm_head`. The tuned lens uses 30 affine translators trained on WikiText-2. Ridge regression initializes each translator against the final residual stream. A sampled KL stage then matches the final distribution over the teacher's top 128 tokens and 128 random negatives. The training report is available in `tuned-lens-training.json`. On eight held-out prompts, the tuned lens reduced mean full-vocabulary KL from 20.51 for the plain logit lens to 2.15. Mean top-10 overlap with the final distribution rose from 1.67 to 3.74 tokens. The Q8 ONNX export preserved at least 8 of the PyTorch lens's top 10 tokens on every tested layer. ## Limitations - The Q4 graph does not reproduce the float32 model exactly. - The graph does not use a KV cache and recomputes the whole sequence when the selected token changes. - Small base models produce weak and sometimes incoherent continuations. - The tuned lens is a course artifact trained on WikiText-2. It is not a research-grade checkpoint or a causal explanation of the model's computation. - Attention weights show values calculated inside the model. They do not establish a causal explanation for an output. - This export is meant for teaching and inspection, not production inference. ## Source The export script and browser laboratory were created for the Engenharia Assistida por IA course. Hugging Face published the original SmolLM2 weights and architecture under the Apache 2.0 license.