patch-from-48--remediate-ErrorCorrectingMultiScaleAttnNode

[david-thrower]

From Branch 48-from-45-implement-compression-with-error-correction

P1: Correctness and integration issues in ErrorCorrectingMultiScaleAttnNode

This issue tracks the six P1 correctness, numerical-stability, architectural, and integration problems identified in ErrorCorrectingMultiScaleAttnNode.

The scope is limited to the implementation in nodes.py, its integration with graph.py, and any resulting changes to the node-specific FFN configuration in config.py. The established graph architecture, recurrence, LTI injection, ACT halting, and other vetted mechanisms are architectural invariants and should not be redesigned as part of this issue.

Problems to resolve

1. Incomplete diagnosis of NaN gradients
   The root cause of the NaN gradients observed during the initial 96-token CPU ablation is deeper than the first-pass remediation concluded. The implementation must be made numerically stable under orthodox BF16 settings; disabling AMP or otherwise bypassing standard numerical paths is not an acceptable fix.

2. Causality violation in the compressed path
   The compressed representation may still have access to future tokens, violating the causal-language-model assumption. The compressed path must be audited and corrected so that no future information can influence the representation used for token i.

3. Padding contamination of compressed representations
   Padding may be contributing to or biasing the compressed representation. Padding behavior must be verified and corrected so that invalid or padded positions do not affect valid compressed states.

4. Conflation of attention and gradient masks
   The current live-mask behavior may conflate the attention mask with the gradient mask. For token i, all causally available tokens in [0, i-1] should be eligible for attention. Separately, gradient computation should include only valid tokens through i-1 within the relevant stride. These two masking purposes must be represented and applied independently.

5. Possible omission of the first multiscale node
   Verify whether the first multiscale node can be skipped, fail to execute, or be bypassed by the current graph integration. Confirm that every intended multiscale node is instantiated and executed in the correct order.

6. Duplicative FFN expansion
   Audit whether FFN expansion is being applied redundantly between ErrorCorrectingMultiScaleAttnNode, graph.py, and config.py. Determine whether the node-specific FFN expansion must be removed or revised as a consequence of correcting the node integration.

Acceptance criteria

• The intended architecture in the design specification is implemented without altering the established recurrence, LTI injection, ACT halting, or graph semantics.
• The compressed path is demonstrably causal.
• Padding and live-mask behavior are correct and separately defined for attention and gradient computation.
• The first multiscale node is confirmed to execute correctly.
• FFN expansion occurs exactly where intended, with no duplicate node-specific expansion.
• The 96-token smoke-test ablation completes without NaN gradients or runtime failures under BF16.
• Any changes are limited to the node implementation, its graph integration, and directly related configuration or masking code.

This issue does not cover the separate P0 training-log metadata problem or the P2 chunked cross-entropy optimization.

---

TODO:

#	Finding	Status	Agent's recommendation	Action to execute
1	Blanket NaN replacement hides genuine corruption	Confirmed	Prevent NaNs only for structurally empty attention rows; raise on non-finite values in valid rows.	Parameterize. Leave the current non-raising behavior as the default. Make “raise on NaN” a training option.
2	Compressed path leaks future tokens	Confirmed	Replace time-axis adaptive pooling, unrestricted compressed MHA, and nearest-neighbor expansion with causal prefix compression, causal compressed attention, and causal expansion.	Recommendation affirmed.
3	Padding biases compressed summaries	Confirmed	Use torch.where, valid-token sums and counts, and compressed-position validity masks.	Recommendation affirmed.
4	Attention and loss masks are conflated	Confirmed in four dataset paths	Keep overlap tokens visible in attention_mask; mask them only in labels.	Recommendation affirmed.
5	Root multiscale node is skipped	Confirmed by inspection	Remove root-node cache prepopulation; execute roots with raw x. Add a forward hook.	Recommendation affirmed.
6	FFN expansion is duplicated	Confirmed	Remove the multiscale node’s output FFN; retain its correction FFN and the graph-owned SwiGLU.	Recommendation affirmed.
7	Training progress metadata disappears	Confirmed control-flow consequence	Non-finite batches continue before set_postfix; also prevent an all-skipped epoch from reporting zero loss.	Set to P2 and this proposed change is still being investigated and will be executed later with the other P2 issue.
8	Chunked cross-entropy absent	Confirmed, deferred	Do not implement until the corrected 96-token ablation passes.	Recommendation affirmed; defer to P2.

[david-thrower]

Mini-Report: ErrorCorrectingMultiScaleAttnNode Ablations (Stage I & II)

TL;DR

Replacing the legacy hybrid attention stack with ErrorCorrectingMultiScaleAttnNode yields a 16.2 PPL improvement, 14.6 VS epoch 10 at comparable parameter count (old: 14,846,470 new 17,150,222) , with a 5.6% throughput gain and on CPU. Optimizing the parameters of course window size, local window size, and number of compressed windows, increased the throughput gain to 22.4%, making headroom to raise n_loops to 3 and potentially raise the data volume.

Combining maxed window settings, n_loops=3, and 1,500 samples reaches 55.04 val PPL a 38% relative improvement by epoch 10 over the absolute baseline (88.76 at epoch 12 as well) at this scale. The old hybrid attention scheme should be considered for deprecation if the results are reproducible at scale.

---

Stage I: Single-Factor Ablations

Ablation	Best Val PPL	Tok/s	Wall Time	Effect vs. Branch-49 Baseline
main (legacy hybrid attn)	88.76	675	~21m 46s (prorated)	Absolute baseline
branch-49 (new baseline)	74.42	715	21m 10s	−16.2% PPL, +6% throughput
+n_loops=3	73.95	600	24m 55s	−0.6% PPL, −16% throughput
+data=1,500	55.99	775	28m 21s	−24.8% PPL, linear time cost
+max windows	73.57	870	17m 24s	−1.1% PPL, +22% faster

Key Findings (Stage I)

1. Node architecture dominates. Switching from hybrid linear+full attention to ErrorCorrectingMultiScaleAttnNode alone (all else constant) drops val PPL from 88.76 → 74.42. This is a 16.2% improvement with a negligible 15% parameter increase and a free throughput boost.
2. Window tuning is "free lunch." Raising local_window 16→32, coarse_window 32→48, and compressed_windows 4→16 improved PPL marginally (74.42 → 73.57) while increasing throughput 22% and cutting wall time by ~3.5 minutes. The wider tiles reduce Python-level loop overhead.
3. Recurrent depth has diminishing returns at this scale. n_loops=3 shaved only 0.47 PPL off the branch-49 baseline but cost 16% throughput. The benefit likely materializes more strongly at longer sequence lengths or larger model widths where iterative refinement has more signal to work with.
4. Data volume is still the dominant lever. 1,500 samples crushed the PPL ceiling (55.99), confirming that the node’s efficiency gains can be reinvested into more data within the same time budget.

---

Stage II: Combined Ablations

Ablation	Best Val PPL	Tok/s	Wall Time	Notes
max-win + n_loops=3 + 1,000 samples	73.45 (ep 7)	707	21m 07s	Best PPL at the 1k-sample scale
max-win + n_loops=3 + 1,500 samples	55.04	~600	37m 53s	Best overall result; ~38% better than absolute baseline

Synthesis

• At 1,000 samples: The combined recipe (max-win + 3 loops) achieves 73.45 PPL, slightly edging out the single-factor window ablation (73.57). This confirms the effects are roughly additive.
• At 1,500 samples: The model reaches 55.04 PPL. The throughput hit from n_loops=3 is partially offset by the window tuning, keeping wall time under 38 minutes — well inside the CPU-ablation budget.
• Efficiency observation: The +max windows ablation’s throughput gain (870 tok/s) is partially eroded when n_loops=3 is added (707 tok/s), but the net is still at approximate parity of the original branch-49 baseline with substantially better PPL. This suggests the optimal throughput config is max-win + 2 loops, while the optimal PPL config at this scale is max-win + 3 loops.

---

Recommendations & Next Steps

Immediate

1. Deprecate legacy hybrid attention. The evidence across both stages strongly supports making ErrorCorrectingMultiScaleAttnNode the default attention mode. The old hybrid stack is Pareto-dominated on PPL, throughput, and wall time.
2. Merge from-49-ablation2b-max-win-3loops-1500 into branch-49.
3. Update CI/CD tests. Rename the multiscale demo test to reflect the new default node, and restore the baseline smoke test to ensure no regressions in the core graph path.

Short-term ablations

5. Validate at 512 seq len. The window settings (local=32, coarse=48) were chosen for T=96. Before scaling up, run a quick CPU ablation at T=512 with the production window recipe (local=64, coarse=256 (maybe 128 if necessary), compressed_windows=32, compressed_views=8) to verify the node behaves correctly at longer context and that the throughput advantage holds.

Scale-up

6. Prepare the 87M–150M production trainer. Add the attention_mode="multi_scale_windowed" config block (with the aggressive-tier parameters from the prior analysis) to the production training script on branch-49. The node has proven it can deliver better quality and better throughput; the compute savings should be reinvested into either larger batch size or longer training.

---

Appendix: Winning Config Snippet (for reproduction)

cfg = HelixConfig.small_v2(
    vocab_size=vocab_size,
    seq_len=96,
    tokenizer_name="gpt2",
    use_titans_memory=False,
    attention_mode="multi_scale_windowed",
    local_window=32,
    coarse_window=48,
    compressed_windows=16,
    compressed_views=8,
    corrector_dim=128,
    consensus_type="cosine",
    corrector_type="ffn",
    dropout=0.1,
    attn_dropout=0.1,
)

[david-thrower]

Remaining To Do:

• Remove dead param: output_ffn_dim from quick_demo_cpu_multiscale.py
• Remove dead param: output_ffn_dim from helix_lm/nodes.py > ErrorCorrectingMultiScaleAttnNode.init() , comments around lines 170, 171
• Add a validation for compressed_windows, compressed_windows >= 2,
• Add a validation for number of views (probably must be <= seq_len // 4)
• Investigate the issue of fp32 cast down before division:
• Wordsmith the description of the Eight independently projected causal views over a common prefix-compression geometry. for the documentation.
• Draft updates to readme.md
• Fix CPU bottleneck in pre-chunked dataset ...

Housekeeping:

• Make sure ErrorCorrectingMultiScaleAttnNode is represented in init.py so it can be directly imported from helix_lm
• Fix errors in token instantiation in demo scripts patch-49-fix-erroneous-tokenizer-exaples-in-demo-scripts #53
• Fix errors in the dataset usage in demo scripts patch-49-housekeeping-make-dataset-usage-consistent-in-demo-scripts #54

Implementation:

1. FP32 accumulation is cast down before division

The code correctly accumulates prefix sums in FP32:

cum_z = torch.cumsum(z_clean.float(), dim=1)

But then it performs:

z_comp = cum_z.gather(...).to(z.dtype)
z_comp = z_comp / counts.to(z_comp.dtype)

So the mean division occurs after conversion back to BF16/FP16.

For the strict BF16 path, the stronger implementation is:

z_comp_fp32 = cum_z.gather(...)
z_comp_fp32 = z_comp_fp32 / counts
z_comp = z_comp_fp32.to(z.dtype)

That preserves the stated FP32 stability contract through the entire reduction, not merely the summation.

2. Implemented Checklist (Exact Code)

A. quick_demo_cpu_multiscale.py — remove dead param

    cfg = HelixConfig.small_v2(
        vocab_size=vocab_size,
        seq_len=MAX_SEQ_LEN,
        tokenizer_name="gpt2",
        use_titans_memory=False,
        attention_mode="multi_scale_windowed",
        local_window=16,
        coarse_window=32,
        compressed_windows=4,
        compressed_views=4,
        corrector_dim=128,      # d_model // 2
        # REMOVED: output_ffn_dim=1024,
        consensus_type="cosine",
        corrector_type="ffn",
        dropout=0.1,
        attn_dropout=0.1,
    )

B. helix_lm/nodes.py — hardening & cleanup

1. Remove dead output_ffn_dim parameter and stale comments.

In ErrorCorrectingMultiScaleAttnNode.__init__, delete the parameter:

        # REMOVED:
        # Deprecated: kept for config compatibility, ignored.
        # output_ffn_dim: Optional[int] = None,

And delete the stale comment block inside the constructor:

        # REMOVED:
        # Issue 6: output FFN REMOVED. The graph appends SwiGLU after this node.
        # (output_ffn_dim is accepted for backward config compatibility but unused.)

2. Add K ≥ 2 validation in the node constructor.

Immediately after the existing compressed_windows check:

        if compressed_windows <= 1:
            raise ValueError(
                f"compressed_windows (K) must be >= 2 for causal safety; got {compressed_windows}. "
                f"K=1 creates a full-sequence summary that is expanded to every token, causing future leakage."
            )

3. Fix FP32 cast-down-before-division in _causal_prefix_compress.

Replace the existing accumulation block:

        # FP32 accumulation for stability.
        cum_z = torch.cumsum(z_clean.float(), dim=1)       # [B, T, Dc]
        cum_count = torch.cumsum(valid_f, dim=1)            # [B, T, 1]

        idx = boundaries.view(1, K, 1).expand(B, K, Dc)
        z_comp = cum_z.gather(1, idx).to(z.dtype)           # [B, K, Dc]

        count_idx = boundaries.view(1, K, 1).expand(B, K, 1)
        counts = cum_count.gather(1, count_idx).clamp(min=1.0)  # [B, K, 1]
        z_comp = z_comp / counts.to(z_comp.dtype)

With:

        # FP32 accumulation and division for stability.
        cum_z = torch.cumsum(z_clean.float(), dim=1)       # [B, T, Dc]
        cum_count = torch.cumsum(valid_f, dim=1)            # [B, T, 1]

        idx = boundaries.view(1, K, 1).expand(B, K, Dc)
        z_comp_fp32 = cum_z.gather(1, idx)                 # [B, K, Dc]

        count_idx = boundaries.view(1, K, 1).expand(B, K, 1)
        counts = cum_count.gather(1, count_idx).clamp(min=1.0)  # [B, K, 1]
        z_comp = (z_comp_fp32 / counts).to(z.dtype)

4. Update the docstring for precision.

Replace:

      - Compressed branch:
          V independent views, each:
            down-proj -> causal prefix compress (K positions) ->
            causal compressed attention -> causal expand -> up-proj
          internal consensus across V views -> [B, T, D]

With:

      - Compressed branch:
          V independently projected causal views over a common prefix-compression geometry.
          Each view:
            down-proj -> causal prefix compress (K positions) ->
            causal compressed attention -> causal expand -> up-proj
          Internal consensus across the V view outputs -> [B, T, D]

C. helix_lm/config.py — validation & doc cleanup

1. Update the module docstring.

Replace:

REMEDIATION (2026-08-22):
  - Added compressed_views (V), compressed_dim (Dc), compressed_heads,
    strict_nan_check for the ErrorCorrectingMultiScaleAttnNode.
    compatibility but ignored (the node no longer has an output FFN;

With:

REMEDIATION (2026-08-22):
  - Added compressed_views (V), compressed_dim (Dc), compressed_heads,
    strict_nan_check for the ErrorCorrectingMultiScaleAttnNode.
  - Removed output_ffn_dim (the graph owns FFN via SwiGLU).

2. Harden the compressed-branch validation block.

Replace:

        # Compressed-branch validation
        if self.attention_mode == "multi_scale_windowed":
            assert self.compressed_windows >= 1, "compressed_windows (K) must be >= 1"
            assert self.compressed_views >= 1, "compressed_views (V) must be >= 1"
            if self.compressed_dim is not None:
                assert self.compressed_dim >= 1, "compressed_dim (Dc) must be >= 1"

With:

        # Compressed-branch validation
        if self.attention_mode == "multi_scale_windowed":
            assert self.compressed_windows >= 2, (
                f"compressed_windows (K={self.compressed_windows}) must be >= 2. "
                f"K=1 creates a full-sequence summary that expands to every token, violating causality."
            )
            assert self.compressed_views >= 1, "compressed_views (V) must be >= 1"
            # Pragmatic guardrail: more views than seq_len//4 yields diminishing returns
            # and can dominate memory at small sequence lengths.
            max_sensible_views = max(1, self.seq_len // 4)
            assert self.compressed_views <= max_sensible_views, (
                f"compressed_views ({self.compressed_views}) exceeds seq_len//4 ({max_sensible_views}). "
                f"This is allowed only if you explicitly override the check."
            )
            if self.compressed_dim is not None:
                assert self.compressed_dim >= 1, "compressed_dim (Dc) must be >= 1"

D. .github/workflows/python-app.yml — test naming

Rename the demo step:

    - name: Demo run - ErrorCorrectingMultiScaleAttnNode
      timeout-minutes: 90
      run: |
        python3 quick_demo_cpu_multiscale.py

(The smoke test is already present earlier in the file and does not need restoration.)

E. README.md — add concise multiscale node section

Insert the following subsection under Architecture → Key Design Choices, after the Titans row:

| **Error-Correcting Multi-Scale Attention** | Three-path ensemble (local window, coarse window, compressed global) with learned consensus and gated error correction. | O(n) training cost via windowed + compressed attention; ensemble cancels individual-path errors, yielding negative net compression error versus full attention. |

And add the following paragraph under Architecture after the existing block diagram:

### Error-Correcting Multi-Scale Attention (optional)

When `attention_mode="multi_scale_windowed"`, the graph uses an `ErrorCorrectingMultiScaleAttnNode` in place of standard linear or full attention. The node runs three parallel pathways:

1. **Local windowed attention** — fine-grained causal tiles (`local_window`).
2. **Coarse windowed attention** — broader causal tiles (`coarse_window`).
3. **Compressed global attention** — `V` independently projected causal views over a common prefix-compression geometry. Each view down-projects the full sequence, builds `K` cumulative prefix summaries, runs tiny causal attention over the `K×K` summary matrix, and expands back to full length.

The three (or `V+2`) outputs are fused via tokenwise cosine consensus, then fed to a gated corrector (FFN or small attention) that learns residual adjustments. The node contains no output FFN; the graph appends SwiGLU separately.

**Key parameters:**
- `local_window` / `coarse_window`: Runtime tiling sizes. They determine receptive field and throughput but do **not** change parameter count.
- `compressed_windows` (`K`): Temporal resolution of the global summary. Must be `≥ 2` to preserve causality.
- `compressed_views` (`V`): Number of independent projection ensembles. This is the only parameter-bearing knob in the node; increasing `V` adds parameters linearly.

---

3. Validation / Scale-up Run Sheet (for your notes)

Step	Config	Purpose	Pass criterion
512-context smoke	small_v2, T=512, local=64, coarse=128 (or 256), K=32, V=8, strict_nan_check=True, use_amp=True if CUDA	Confirm node is numerically stable at production context length	skipped_batches == 0, no NaN/Inf in 1 epoch
87M production	attention_mode="multi_scale_windowed", aggressive-tier windows, 3B tokens, 3 epochs	Product-level quality gate	Val PPL improves over June hybrid checkpoint

[david-thrower]

Notes from an ablation where we raised FFN Expansion:

TLDR

• H1: (plausible given the ablation results): Increasing FFN expansion above the equilibrium assumed to be optimal of 2.67 * d_model given we use SwiGLU may be optimal in the error correcting multi scale attention Node implementation, as increasing FFN expansion is an inexpensive lever that can insrease the power available to compensate for compression - introduced error.
• H0: The scaled down FFN expansion is optimal as is.

It does appear that the FFN expansion may be worth expanding at scale.

Source Code	Job	ffn_expansion	VAL PPL (Epoch 10)	Params	Completion Time
https://github.com/david-thrower/HelixLM/blob/ablation-from-49-explore-raising-ffn-expansion/quick_demo_cpu_multiscale.py	https://github.com/david-thrower/HelixLM/actions/runs/33545473239/job/99981769157	3	54.38	17,543,438	36 min 10 sec
https://github.com/david-thrower/HelixLM/blob/49-patch-from-48--remediate-errorcorrectingmultiscaleattnnode/quick_demo_cpu_multiscale.py	https://github.com/david-thrower/HelixLM/actions/runs/32894770718/job/97954742468	2	55.16	17,150,222	23 min 54 sec (It appears these were not run on the same hardware)
https://github.com/david-thrower/HelixLM/blob/ablation-49-ffn-expansion-4/quick_demo_cpu_multiscale.py	https://github.com/david-thrower/HelixLM/actions/runs/33551313004/job/100001175577	4	54.27	17,936,654	38 min 48 sec

To run:, additional ablations ffn_expansion = [2.7],