How ICE-ICL Works#

A visual walkthrough of the ICL backbone — from the nanotabicl baseline through each architectural innovation. Written for readers who want to understand what was redesigned and why, without reading the source code.

Where the Work Is#

ICE-ICL is an ICL backbone — the component responsible for turning a set of labeled training rows plus unlabeled test rows into predictions. Everything else is untouched: the structural causal model prior that generates synthetic training data, the column and row preprocessing, the output projection. All architectural innovations are concentrated inside the ICL backbone, specifically in how the attention mechanism within it works.

The shared preprocessing stack is identical across all twelve arms (A–D9):

flowchart TD
  subgraph FRONT["Shared Frontend — identical across all arms"]
    IN["x  ·  B × N × C"] --> STD["Standardise\n(train-row statistics)"]
    STD --> GRP["Cyclic feature grouping\noffsets 0, 1, 3  →  B × N × C × 3"]
    GRP --> PROJ["Linear 3 → 128\nx_embed  ·  B × N × C × 128"]
    PROJ --> COL["Column Attention × 3\nInducedTransformerBlock\nk, v = train rows only"]
    COL --> ROW["Row Attention × 3\nTransformerBlock + RoPE\n+ 4 CLS tokens per row"]
    ROW --> CLS["Flatten CLS\nr_i ∈ ℝ⁵¹²  per row"]
  end
  subgraph ICL["ICL Backbone — all architectural innovations are here"]
    ARMS["Arms A · B · C · D–D9\nsee sections below"]
  end
  CLS --> ICL
  ICL --> OUT["logits  ·  B × N_test × n_classes"]

The column attention blocks propagate cross-feature interactions within each row. The row attention blocks integrate cross-row context for each feature. Four CLS columns per row aggregate the result into a single \(\mathbf{r}_i \in \mathbb{R}^{512}\) — the input the ICL backbone receives.

Arm A — The Baseline#

Arm A is the nanotabicl backbone: 12 identical TransformerBlocks with QASSMax attention scaling. It is a self-contained minimal reproduction of the TabICLv2 ICL stack, trained from scratch on SCM priors (27.6 M parameters).

Layers 1–11 process all rows together with train rows as keys/values. Layer 12 uses test queries against train keys to produce predictions:

Layers 1–11:  q = [train ‖ test],  k, v = train only  →  update all rows
Layer 12:     q = test only,        k, v = train only  →  logits

The attention scaling at each layer uses QASSMax:

\[\begin{split}\tilde{q}_{h,i} = \underbrace{W_{\text{base}}(\log n)}_{{\substack{\text{population gain} \\ \text{one vector per head, dim} \\ \text{function of } n \text{ only}}}} \;\odot\; \underbrace{\bigl(1 + \tanh\!\bigl(\mathrm{MLP}_q(q_{h,i})\bigr)\bigr)}_{\substack{\text{per-query modulation} \\ \in (0,\,2) \\ \text{initialised at } 1}} \;\odot\; q_{h,i}\end{split}\]

where \(n\) is the number of training rows and \(h\) indexes attention heads. Scaling \(q\) before the inner product \(q k^\top / \sqrt{d}\) is equivalent to multiplying the inverse softmax temperature by the same factor — sharper attention for larger scale, softer for smaller.

The critical limitation: \(W_{\text{base}}(\log n)\) is identical for every dataset of size \(n\). Two batches — one with cleanly separated classes, one with random labels — receive the same bandwidth because the input to \(W_{\text{base}}\) is only the scalar count \(\log n\), with no information about data geometry.

Arm A → Arm B: Separating Estimation from Decoding#

In Arm A, test rows appear as queries in all 12 ICL layers — including the 11 layers where the model is building the train representations that will serve as keys for prediction. Test queries participate in computation that should be about estimating the label-generating function, not yet about decoding.

Arm B eliminates this structural mixing. Train embeddings are frozen after the frontend and never updated by the ICL backbone. Test rows run 12 dedicated cross-attention passes against those fixed keys:

flowchart LR
  subgraph A["Arm A  ·  12 self-attention blocks"]
    direction TB
    A_IN["all rows  train + test"] --> A1["ICL block 1\nq = all rows\nkv = train\nupdate ALL rows"]
    A1 --> A2["ICL block 2\nq = all rows\nkv = train\nupdate ALL rows"]
    A2 --> ADOTS["· · · ×9 · · ·"]
    ADOTS --> A12["ICL block 12\nq = test  kv = train"]
    A12 --> A_OUT["logits"]
  end
  subgraph B["Arm B  ·  12 cross-attention blocks    +3.8 pp"]
    direction TB
    B_TR["train\nfrozen after frontend"] --> B_KV["stable kv\nfor all 12 layers"]
    B_TE["test\nevolved independently"] --> B1["cross-attn 1\nq = test  kv = frozen train\nonly test updates"]
    B_KV --> B1
    B1 --> B2["cross-attn 2"]
    B2 --> BDOTS["· · · ×10 · · ·"]
    BDOTS --> B_OUT["logits"]
  end

The gain (+3.8 pp, p ≈ 3 × 10⁻⁵²) comes from three structural properties of Arm B:

  1. Frozen train keys — train representations are stable across all 12 decoding layers, providing a consistent key/value source rather than one that is still evolving under self-attention updates.

  2. Dedicated architecture — CrossAttnBlock uses separate query and key/value projections, not the shared in-projection of a self-attention block. Each projection learns a specialised role.

  3. Pure decoding computation — all 12 layers contribute only to decoding test rows, with no capacity shared with train self-refinement.

Arm C — ICE: Closed-Loop Bandwidth Control#

Arm B still inherits QASSMax’s open-loop bandwidth. The attention temperature depends only on \(\log n\) — the same for any two datasets of equal size, regardless of their geometry. Arm C replaces the entire 12-layer ICL backbone with a three-component pipeline:

flowchart TD
  TR_IN["Train rows  r_i ∈ ℝ⁵¹²"] --> SSM1["SSMBlock 1\nbidirectional selective scan\ngated state update over train sequence"]
  SSM1 --> SSM2["SSMBlock 2\nbidirectional selective scan"]
  SSM2 --> POOL_KV["per-row hidden states  h_i"]
  SSM2 --> POOL_Q["learned pooling query"]
  POOL_KV --> POOL_A["Attention pooling"]
  POOL_Q --> POOL_A
  POOL_A --> HG["h_global ∈ ℝ⁶⁴\nchannel state\ncluster structure · noise level\neffective dimensionality"]
  HG --> TAU["Shared MLP → n_heads outputs\nτ_k = 1 + tanh(MLP(h_global))_k\nτ_k ∈ (0, 2)  init 1"]
  SSM2 --> ENRICH["key enrichment\ntrain kv + SSM hidden state"]
  TAU --> XA1["ConditionedCrossAttn 1\nq = test · q scaled by τ_k\nkv = enriched train"]
  ENRICH --> XA1
  XA1 --> XA2["ConditionedCrossAttn 2"]
  XA2 --> XA3["ConditionedCrossAttn 3"]
  XA3 --> XA4["ConditionedCrossAttn 4"]
  XA4 --> OUT["logits"]
  TE_IN["Test rows  r_j ∈ ℝ⁵¹²"] --> XA1

Identification (SSM scan). Two bidirectional SSM blocks process the training sequence with gated state updates:

\[h_t = \exp(A \cdot \mathrm{gate}_t) \odot h_{t-1} + B_t \cdot x_t\]

Each position integrates context from the full training sequence in O(n) time. \(A\) is a learned constant; \(B_t\) and \(\mathrm{gate}_t\) are data-dependent projections of the input. The exponential decay term \(\exp(A \cdot \mathrm{gate}_t)\) controls how much history is retained — the model learns to selectively reset state at class boundaries or distribution shifts.

Conditioning (attention pooling → τ). A learned query vector attends over all SSM hidden states to produce a single global summary:

\[h_{\text{global}} \in \mathbb{R}^{64}\]

This 64-dimensional bottleneck encodes dataset-level geometry — cluster count, class separability, noise level, effective dimensionality. The bottleneck forces the model to distil distributional structure rather than memorise individual points.

Per-head attention bandwidths are derived from this geometry via a single shared MLP with \(n_\text{heads}\) outputs:

\[[\tau_1, \ldots, \tau_K] = 1 + \tanh\!\bigl(\mathrm{MLP}(h_{\text{global}})\bigr), \quad \tau_k \in (0,\,2)\]

Initialised at \(\tau_k = 1\) (neutral). During training each head’s output dimension learns to request more or less bandwidth based on the geometry it specialises in: a head capturing local decision boundaries learns a narrow kernel (high τ); a head capturing global trends learns a wide kernel (low τ). Bandwidth is now closed-loop — derived from the actual data distribution, not from \(\log n\).

Equalization (4× conditioned cross-attention). Test queries attend to enriched train keys, with queries scaled by the per-head \(\tau_k\):

\[\tilde{q}_{h,j} = \tau_k \cdot q_{h,j}\]

No test-to-test attention at any point. Each test row is decoded independently.

Arm C uses 12.6 M parameters — 2.2× fewer than Arm A — while outperforming it on all evaluation metrics. The gain comes from structural specialisation: 2 SSM blocks for estimation and 4 cross-attention blocks for equalization replace 12 uniform self-attention blocks doing all jobs simultaneously.

Arm D — ICED: Decision Feedback#

Arm D adds one round of turbo-style feedback on top of Arm C. First-pass soft predictions are fed back as pseudo-pilots to refine the channel estimate:

  1. Pass 1: Arm C pipeline → soft predictions \(\hat{p}_j = \mathrm{softmax}(\mathrm{logits}_1[j])\).

  2. Re-estimation: \(\hat{p}_j\) projected into the embedding space and concatenated with train rows. The SSM re-scans \([\mathrm{train} \| \mathrm{pseudo\text{-}test}]\) → updated \(h_{\text{global},2}\), \(\tau_2\), enriched keys₂.

  3. Gated blend:

    \[(\tau,\,\mathrm{enr}) = g \cdot (\tau_2,\,\mathrm{enr}_2) + (1-g) \cdot (\tau_1,\,\mathrm{enr}_1), \quad g = \sigma(\mathrm{gate})\]

    Gate scalar initialised at \(-3\) → \(g \approx 0.05\) at start. Arm D begins as pure Arm C and opens the feedback gate only when the second-pass conditioning consistently improves predictions.

  4. Pass 2: shared equalizer with blended conditioning → final logits.

All heavy components are shared between passes (SSM, cross-attention, output MLP). Arm D adds 5,633 parameters (~0.04% overhead) over Arm C.

Arm Comparison#

Arm

Params

ICL backbone

A · Baseline

27.6 M

12× self-attention + QASSMax. All rows in q (layers 1–11), only train in kv. Bandwidth from \(\log n\) only.

B

27.6 M

12× dedicated cross-attention + QASSMax. Train frozen, test does 12 pure decoding passes. +3.8 pp vs A.

C · ICE

12.6 M

2× BiSSM → \(h_{\text{global}} \in \mathbb{R}^{64}\) → \(\tau_k\) per head + enriched keys → 4× conditioned cross-attention. Closed-loop bandwidth. +6.7 pp vs A. 2.2× fewer parameters.

D · ICED

≈ 12.6 M

Arm C + gated turbo feedback. Gate init \(\sigma(-3) \approx 0.05\). +6.6 pp vs A. D3 (incremental) leads on dataset wins.

Nine additional variants (D2–D9) isolate specific design choices within decision feedback and channel estimation — see Ablation Design for definitions and results. Arms D6–D9 are motivated by Ljung, System Identification (1999).

For full implementation details and design rationale see Architecture and Why Redesign the ICL Backbone?. For experimental results see Research.