Attention Pattern

Back to design

This note tracks the Apr22 attention-pattern analysis for adaptation_mode=full and loss setting beta0p01_d50_lin50. The analysis asks whether generated payload tokens with high z-free base-model surprisal also place more latent-conditioned attention on the continuous latent placeholder <z_empty>.

Full CUDA analysis completed on 2026-04-25 with 256 generated test traces for each of six tasks. The final artifact set contains 1536 traces, 1535 analyzable traces, and 437296 token-layer rows. One sorting_algorithms trace produced no aligned payload sequence and is retained in trace metadata with analyzable=false.

Method

The base-model comparison is base_token_nll = -log p_base(y_t | x, y_<t), the token-level base-model surprisal of the same generated payload token under the frozen z-free reference model after removing the latent span. Most generated tokens have surprisal extremely close to zero, with a small long tail, so the scatter view uses empirical base-model surprisal percentile on the x-axis. This preserves the rank ordering of surprising tokens while making the dense low-surprisal bulk readable.

The example bar panels use raw base-model surprisal so the units match -log p_base(y_t | x, y_<t) directly.

The primary token metric is predictive_z_attention: the maximum, over heads, of attention from the query position whose logits predict payload token y_t to the <z_empty> key position. The secondary metric is realized_token_z_attention: the same max-over-heads score from the realized payload token position to <z_empty>.

Hard-attention diagnostics use the full key-wise attention replay. For each payload token, layer, and head, the analysis records whether <z_empty> is the top-attended valid causal key, its rank among valid keys, and its margin over the best non-z key. Ties within 1e-12 count as top-attention. The local hard-attention diagnostic uses predictive_z_top_head_fraction as an any-head indicator: 1 when at least one head makes <z_empty> the top valid key at the prediction query for y_t, and 0 otherwise.

Outputs

• Trace metadata: path:experiments/synthetic_sequences/disentangled/analysis/aggregates/continuous_cvae-Apr22-beta-schedule-followup_20260422/data/attention_pattern_traces.jsonl
• Token/layer table: path:experiments/synthetic_sequences/disentangled/analysis/aggregates/continuous_cvae-Apr22-beta-schedule-followup_20260422/data/attention_pattern_token_layer.csv
• Layer summary: path:experiments/synthetic_sequences/disentangled/analysis/aggregates/continuous_cvae-Apr22-beta-schedule-followup_20260422/data/attention_pattern_layer_summary.csv
• Bootstrap layer summary: path:experiments/synthetic_sequences/disentangled/analysis/aggregates/continuous_cvae-Apr22-beta-schedule-followup_20260422/data/attention_pattern_bootstrap_layer_summary.csv
• Hard-attention layer summary: path:experiments/synthetic_sequences/disentangled/analysis/aggregates/continuous_cvae-Apr22-beta-schedule-followup_20260422/data/attention_pattern_hard_layer_summary.csv
• Head summary: path:experiments/synthetic_sequences/disentangled/analysis/aggregates/continuous_cvae-Apr22-beta-schedule-followup_20260422/data/attention_pattern_head_summary.csv
• Hard-attention head summary: path:experiments/synthetic_sequences/disentangled/analysis/aggregates/continuous_cvae-Apr22-beta-schedule-followup_20260422/data/attention_pattern_hard_head_summary.csv

Figures

Per-task example panels are copied under figures/attention_pattern_examples_by_task_panels/. These local token-level views use average-over-head soft attention to <z_empty>, so a cell shows the layer's mean attention-to-z across heads for that token prediction query.

Per-task hard-attention example panels are copied under figures/attention_pattern_hard_examples_by_task_panels/. These use the same traces and token rows as the soft example panels, but the heatmap value is the any-head version of predictive_z_top_head_fraction. A bright cell means that at least one head in that layer makes <z_empty> the top-attended valid key for the token prediction query.

The gray hexbin background is token density within each layer. The x-axis is empirical base-model surprisal percentile, so the right side of each panel contains the most surprising generated tokens. Each colored line is the per-task least-squares line of best fit between surprisal percentile and attention to <z_empty>.

The alignment figures show task-wise patterns only. Pearson correlation measures linear association in the raw values, so it is sensitive to the magnitude of rare high-surprisal tokens and answers whether attention grows approximately linearly with surprisal. Spearman correlation measures rank association, so it asks whether more surprising tokens generally receive more attention regardless of the exact scale. For this exploratory question, Spearman is usually the more robust summary because base-model surprisal is highly skewed; Pearson is still useful for detecting whether the largest-surprisal branching tokens dominate the effect. The prior all-task aggregate line and confidence band were removed because they hid the task heterogeneity.

The heatmap cells show Spearman correlations numerically. A small * marks task/layer cells whose trace-cluster bootstrap 95% interval excludes zero.

Validation

The focused regression tests cover payload token alignment, latent-span removal, <z_empty> indexing after padding, hard-attention top-key detection, tie handling, rank/margin masking, and uniform causal baselines. The test command passed:

source scripts/uv_venv.sh && PYTHONPATH=. uv run pytest experiments/synthetic_sequences/disentangled/tests/test_attention_pattern_analysis.py

Full artifact validation:

• 1536 traces were written, matching six tasks times 256 traces per task.
• 1535 traces have analyzable payload-token rows.
• The trace set covers base_conversion, grid_pathfinding, linear_equation_solving, list_summation, multidigit_addition, and sorting_algorithms.
• The token-layer table has 437296 rows and 1535 unique analyzable trace ids.
• No analyzable payload-token row has missing surprisal, baseline, enrichment, rank, margin, or hard-attention values.
• Soft attention values are bounded in [0, 1].
• Hard-attention any-head indicators are bounded in [0, 1].
• Bootstrap and hard-attention summaries contain all eight layers and all six tasks plus the all-task aggregate.
• V3 rebuild produced split Pearson/Spearman figures, six per-task soft example panels, and six per-task hard example panels.

Main Readout

At the all-task aggregate level, Pearson alignment between base-model surprisal and predictive_z_attention is positive in every layer. Using raw surprisal, Pearson correlations are about 0.30 to 0.34; using log(1 + surprisal), Pearson correlations are stronger, about 0.52 to 0.62. Spearman alignment is also positive in every layer, with the strongest aggregate rank alignment in displayed layer 8.

Aggregate layer profile:

The task-level pattern is heterogeneous. multidigit_addition and grid_pathfinding have the strongest raw-Pearson best layers. sorting_algorithms has weak Pearson but a clearer rank-alignment signal in middle layers. list_summation has its strongest Pearson at displayed layer 4, which is also the layer with high aggregate enrichment.

Best Pearson layer by task:

The hard-attention example panels are the local diagnostic for the same surprisal hypothesis. They ask whether the token rows with high base-model surprisal are also the rows where <z_empty> becomes the top key in at least one head. This is stricter than soft attention: a token/layer can assign nontrivial probability mass to z without making z the maximum-attended key, while the hard indicator records whether z wins the head-level competition anywhere in the layer.

The strongest head-level raw-Pearson alignment is displayed layer 4, head 6 (0.310). The highest hard-top-key head is displayed layer 3, head 3, where <z_empty> is the top valid key for 0.144 of token queries.

Top head-level Pearson alignments:

Run Command

The V3 figure pass reused the existing full-run CSVs:

source scripts/uv_venv.sh
PYTHONUNBUFFERED=1 PYTHONPATH=. python -u experiments/synthetic_sequences/disentangled/analysis/run_continuous_cvae_apr22_attention_pattern.py --rebuild-from-existing