README.md

---
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.