NLP probing

Requires the [nlp] extra (pip install igl[nlp]).

igl.nlp

Language-model probing helpers. Requires the [nlp] extra.

Verified plumbing for head-in-the-loop measurements on HuggingFace causal LMs: :func:resolve_head probes the head convention instead of assuming it, :class:ContinueFrom builds identity-gated patch-and-continue closures, :func:extract_activations produces disjoint fit/eval splits, and :func:requirement_dimension sweeps projection arms against the frozen head. Importing any submodule raises :class:igl.IGLDependencyError when transformers is missing, so the failure is loud and actionable.

Projector = Callable[[torch.Tensor, int], torch.Tensor] module-attribute

A projection arm: (states [N, C], k) -> reconstructed states [N, C].

ActivationSet dataclass

Extracted hidden states, split into disjoint fit and eval sets.

Attributes:

Name	Type	Description
fit	Tensor	Subsampled states [n_fit, C] from the post-eval chunks.
eval_batches	list[tuple[Tensor, Tensor]]	One (states [ctx-1, C], targets [ctx-1]) pair per evaluation chunk; targets are the next-token ids, ready for :func:igl.nlp.requirement_dimension.
layer	int	The zero-based block index the states come out of.
ctx	int	The context length used for chunking.

ContinueFrom

Run blocks layer+1 .. end plus final norm and head on given states.

Parameters:

Name	Type	Description	Default
model	object	A HuggingFace causal LM (GPT-2, LLaMA/Qwen, GPT-NeoX families: block lists at transformer.h, model.layers, or gpt_neox.layers).	required
layer	int	Zero-based block index the states come out of; the closure applies blocks layer+1 onward. Must satisfy 0 <= layer < n_blocks.	required
probe_tokens	int	Length of the identity-gate forward.	16
atol	float	Absolute logits tolerance of the identity gate, in fp32.	0.01

Raises:

Type	Description
IGLConfigError	When the architecture is not recognized, layer is out of range, or the identity gate fails (continuing from the model's own states does not reproduce its logits).

__call__(states)

Continue from states [T, C] to fp32 logits [T, V].

HeadClosure dataclass

A verified hidden-states-to-logits closure.

Attributes:

Name	Type	Description
post_norm	bool	Whether the final hidden state is already normalized (the unembedding applies to it directly).
lm_head	Module	The unembedding module.
norm	Module | None	The final normalization module, when post_norm is false.

__call__(states)

Map states [..., C] to fp32 logits, casting through the head's dtype.

RequirementReport dataclass

Result of a required-dimension sweep.

Attributes:

Name	Type	Description
intact_perplexity	float	Perplexity of the head reading the raw states.
curves	dict[str, dict[int, float]]	Per arm, the perplexity at each swept k.
required	dict[str, dict[float, int | None]]	Per arm, the smallest swept k whose perplexity is within tolerance * intact_perplexity, or None when no swept k qualifies (the requirement exceeds the grid).

extract_activations(model, input_ids, *, layer, ctx=512, n_eval_chunks=20, n_fit=20000, seed=42)

Extract layer activations with disjoint fit and eval splits.

The token stream is cut into ctx-length chunks. The first n_eval_chunks become the evaluation set; the fit pool is sampled (without replacement, seeded) from all positions of the remaining chunks, so fit and eval tokens never overlap.

Parameters:

Name	Type	Description	Default
model	object	A HuggingFace causal LM with hidden-state outputs.	required
input_ids	Tensor	A 1-D token-id tensor.	required
layer	int	Zero-based block index; states are the residual stream after block layer (hidden_states[layer + 1]).	required
ctx	int	Context length per forward.	512
n_eval_chunks	int	Number of leading chunks reserved for evaluation.	20
n_fit	int	Number of fit positions to sample; capped at the pool size.	20000
seed	int	Seed of the fit subsampling.	42

Returns:

Name	Type	Description
An	ActivationSet

Raises:

Type	Description
IGLConfigError	When input_ids is not 1-D or holds fewer than (n_eval_chunks + 1) * ctx tokens (no fit chunk would remain).

require_transformers()

Raise :class:igl.IGLDependencyError unless the nlp extra is installed.

requirement_dimension(projectors, head, batches, *, grid, tolerances=_DEFAULT_TOLERANCES)

Sweep projection arms over k and report per-tolerance requirements.

Parameters:

Name	Type	Description	Default
projectors	Mapping[str, Projector]	Named arms mapping (states, k) to reconstructed states of the same width (e.g. an IGL chart's project-then-reconstruct, or a PCA baseline).	required
head	Callable[[Tensor], Tensor]	The frozen consumer mapping states to logits: an :class:igl.nlp.HeadClosure for the final layer, or an :class:igl.nlp.ContinueFrom for patch-and-continue profiles.	required
batches	Iterable[tuple[Tensor, Tensor]]	(states [N, C], targets [N]) evaluation pairs; targets are next-token ids. Materialized once and reused per (arm, k).	required
grid	Sequence[int]	The k values to sweep, in increasing order.	required
tolerances	Sequence[float]	Multiplicative perplexity tolerances; for each, the report records the smallest swept k within tolerance * intact_perplexity.	_DEFAULT_TOLERANCES

Returns:

Name	Type	Description
A	RequirementReport
	RequirementReport	perplexity curve per arm, and per-tolerance required dimensions.

Raises:

Type	Description
IGLConfigError	When grid is empty, not increasing, or batches is empty.

resolve_head(model, *, probe_tokens=8, atol=0.01)

Probe and verify how the unembedding reads the final hidden state.

Runs a short forward with hidden states, then checks which convention reproduces the model's own logits: the head on the raw last hidden state, or the head on the final norm of it (candidate norms: model.norm, transformer.ln_f, gpt_neox.final_layer_norm).

Parameters:

Name	Type	Description	Default
model	object	A HuggingFace causal LM exposing lm_head (or embed_out) and output_hidden_states.	required
probe_tokens	int	Length of the verification forward.	8
atol	float	Absolute tolerance of the logits comparison, in fp32.	0.01

Returns:

Type	Description
HeadClosure	A verified :class:HeadClosure.

Raises:

Type	Description
IGLConfigError	When no unembedding module is found, or neither convention reproduces the model's logits.

igl.nlp.resolve_head(model, *, probe_tokens=8, atol=0.01)

Probe and verify how the unembedding reads the final hidden state.

Runs a short forward with hidden states, then checks which convention reproduces the model's own logits: the head on the raw last hidden state, or the head on the final norm of it (candidate norms: model.norm, transformer.ln_f, gpt_neox.final_layer_norm).

Parameters:

Name	Type	Description	Default
model	object	A HuggingFace causal LM exposing lm_head (or embed_out) and output_hidden_states.	required
probe_tokens	int	Length of the verification forward.	8
atol	float	Absolute tolerance of the logits comparison, in fp32.	0.01

Returns:

Type	Description
HeadClosure	A verified :class:HeadClosure.

Raises:

Type	Description
IGLConfigError	When no unembedding module is found, or neither convention reproduces the model's logits.

igl.nlp.HeadClosure dataclass

A verified hidden-states-to-logits closure.

Attributes:

Name	Type	Description
post_norm	bool	Whether the final hidden state is already normalized (the unembedding applies to it directly).
lm_head	Module	The unembedding module.
norm	Module | None	The final normalization module, when post_norm is false.

__call__(states)

Map states [..., C] to fp32 logits, casting through the head's dtype.

igl.nlp.ContinueFrom

Run blocks layer+1 .. end plus final norm and head on given states.

Parameters:

Name	Type	Description	Default
model	object	A HuggingFace causal LM (GPT-2, LLaMA/Qwen, GPT-NeoX families: block lists at transformer.h, model.layers, or gpt_neox.layers).	required
layer	int	Zero-based block index the states come out of; the closure applies blocks layer+1 onward. Must satisfy 0 <= layer < n_blocks.	required
probe_tokens	int	Length of the identity-gate forward.	16
atol	float	Absolute logits tolerance of the identity gate, in fp32.	0.01

Raises:

Type	Description
IGLConfigError	When the architecture is not recognized, layer is out of range, or the identity gate fails (continuing from the model's own states does not reproduce its logits).

__call__(states)

Continue from states [T, C] to fp32 logits [T, V].

igl.nlp.extract_activations(model, input_ids, *, layer, ctx=512, n_eval_chunks=20, n_fit=20000, seed=42)

Extract layer activations with disjoint fit and eval splits.

The token stream is cut into ctx-length chunks. The first n_eval_chunks become the evaluation set; the fit pool is sampled (without replacement, seeded) from all positions of the remaining chunks, so fit and eval tokens never overlap.

Parameters:

Name	Type	Description	Default
model	object	A HuggingFace causal LM with hidden-state outputs.	required
input_ids	Tensor	A 1-D token-id tensor.	required
layer	int	Zero-based block index; states are the residual stream after block layer (hidden_states[layer + 1]).	required
ctx	int	Context length per forward.	512
n_eval_chunks	int	Number of leading chunks reserved for evaluation.	20
n_fit	int	Number of fit positions to sample; capped at the pool size.	20000
seed	int	Seed of the fit subsampling.	42

Returns:

Name	Type	Description
An	ActivationSet

Raises:

Type	Description
IGLConfigError	When input_ids is not 1-D or holds fewer than (n_eval_chunks + 1) * ctx tokens (no fit chunk would remain).

igl.nlp.ActivationSet dataclass

Extracted hidden states, split into disjoint fit and eval sets.

Attributes:

Name	Type	Description
fit	Tensor	Subsampled states [n_fit, C] from the post-eval chunks.
eval_batches	list[tuple[Tensor, Tensor]]	One (states [ctx-1, C], targets [ctx-1]) pair per evaluation chunk; targets are the next-token ids, ready for :func:igl.nlp.requirement_dimension.
layer	int	The zero-based block index the states come out of.
ctx	int	The context length used for chunking.

igl.nlp.requirement_dimension(projectors, head, batches, *, grid, tolerances=_DEFAULT_TOLERANCES)

Sweep projection arms over k and report per-tolerance requirements.

Parameters:

Name	Type	Description	Default
projectors	Mapping[str, Projector]	Named arms mapping (states, k) to reconstructed states of the same width (e.g. an IGL chart's project-then-reconstruct, or a PCA baseline).	required
head	Callable[[Tensor], Tensor]	The frozen consumer mapping states to logits: an :class:igl.nlp.HeadClosure for the final layer, or an :class:igl.nlp.ContinueFrom for patch-and-continue profiles.	required
batches	Iterable[tuple[Tensor, Tensor]]	(states [N, C], targets [N]) evaluation pairs; targets are next-token ids. Materialized once and reused per (arm, k).	required
grid	Sequence[int]	The k values to sweep, in increasing order.	required
tolerances	Sequence[float]	Multiplicative perplexity tolerances; for each, the report records the smallest swept k within tolerance * intact_perplexity.	_DEFAULT_TOLERANCES

Returns:

Name	Type	Description
A	RequirementReport
	RequirementReport	perplexity curve per arm, and per-tolerance required dimensions.

Raises:

Type	Description
IGLConfigError	When grid is empty, not increasing, or batches is empty.

igl.nlp.RequirementReport dataclass

Result of a required-dimension sweep.

Attributes:

Name	Type	Description
intact_perplexity	float	Perplexity of the head reading the raw states.
curves	dict[str, dict[int, float]]	Per arm, the perplexity at each swept k.
required	dict[str, dict[float, int | None]]	Per arm, the smallest swept k whose perplexity is within tolerance * intact_perplexity, or None when no swept k qualifies (the requirement exceeds the grid).