attnres-implementation

name	attnres-implementation
description	Attention Residuals (AttnRes) implementation in the PDLI experiment. Covers AttnResDepth, AttnResBlock, and the L2RTP model that uses them. Essential for anyone modifying the transformer architecture or residual connections.

AttnRes (Attention Residuals) Implementation

Mental Model

AttnRes replaces fixed skip-connections with learned softmax attention over all preceding layer outputs. At layer l, a lightweight query/key projection computes input-dependent softmax weights over {h_0, …, h_{l-1}}, and the weighted sum becomes the residual base. This prevents shallow-layer dilution at depth and allows content-dependent routing through depth.

Coverage

Documented: AttnResDepth, AttnResBlock, L2RTP model, Config, file locations Not yet documented: None known Last extended: 2026-06-21

Key Findings

1. Three copies of the same implementation exist

• Canonical: resources/techniques/chen2026attnres.py (149 lines)
• Alternate: resources/techniques/attn_res.py (145 lines)
• LTP copy: notebooks/ye_ole_LTP/techniques/attn_res.py (145 lines)
• All three are functionally identical. The LTP copy is the one actively imported by l2_rtp.py.

2. AttnResDepth — the core mechanism

• Location: All three attn_res.py/chen2026attnres.py files, class AttnResDepth
• What: Computes a depth-attention weighted sum over past layer outputs
• Forward: q_proj(current_input), k_proj(each_past_output) → softmax(q·k/√d) → weighted sum
• Parameters: d_model → depth_head_dim (default 64) via two bias-free Linear layers
• Zero-shot behavior: With zero-init queries, attention starts uniform

3. AttnResBlock — the transformer block

• Location: Same files, class AttnResBlock
• What: PreNorm transformer block using AttnRes for the residual connection
• Forward: h = attn_res(x, past_outputs) + SelfAttn(LayerNorm(h)) then h + FFN(LayerNorm(h))
• Layer 0: Standard residual (no past outputs), so layer 0 uses identity skip
• Bidirectional: No causal mask — all tokens attend to all tokens

4. L2RTP — Recursive Thought Processor (the model using AttnRes)

• Location: notebooks/ye_ole_LTP/model/l2_rtp.py
• Architecture: 4 AttnResBlock layers at d_model=4096, with recurrence
• Recurrence: Concatenates bundle_0 + previous state + STOP token, runs through blocks, extracts refined state from second half of output
• Stop signal: Learnable STOP token + StopMLP (Linear→GELU→Linear→Sigmoid)
• k_max=8 (configurable), stop_threshold=0.5

5. No DAR-style timestep-adaptive routing in code

• The design doc (DD-multi-teacher-distillation-attnres-diffusion.md) describes DAR-style routing but this has NOT been implemented
• The current AttnRes is the basic version — query/key are simple linear projections of current/past hidden states, with no timestep conditioning

6. Config settings for L2

• Location: notebooks/ye_ole_LTP/config.py
• l2_layers=4, l2_nhead=32, l2_ffn_dim=8192, l2_depth_head_dim=64, k_max=8, stop_threshold=0.5

7. Pipeline integration

• LatentThoughtPipeline (pipeline.py) calls L2RTP in Phase A only
• L2 output concatenated with L1 output before passing to L3
• L2 bypassed in Phase 0, Phase B, and Phase Pre

Critical Invariants

1. AttnResBlock layer 0 uses identity skip (no attn_res module created for layer_idx==0)
2. All AttnResBlocks are bidirectional — no causal mask
3. The past_outputs list accumulates sequentially through blocks during a single forward pass
4. L2 always returns same shape as input: [B, N, d_thought]
5. Pipeline concatenation cat([l1_out, l2_out]) happens outside L2 in pipeline.py

Sources

• resources/techniques/chen2026attnres.py — canonical implementation
• resources/techniques/attn_res.py — alternate implementation
• notebooks/ye_ole_LTP/techniques/attn_res.py — LTP copy
• notebooks/ye_ole_LTP/model/l2_rtp.py — L2RTP model
• notebooks/ye_ole_LTP/model/pipeline.py — pipeline integration
• notebooks/ye_ole_LTP/config.py — configuration
• docs/upstream/papers/chen2026attnres.md — paper summary