cs.LGOct 1, 2026

Pooling Helps, Learned Weighting Hurts In-Context: Decomposing Group Attention

Authors: Michael Fore, James Mason Inder, Mrishika Nair, Praneetha Vaddamanu, Sharlina Keshava

Organizations: Amazon

Read PDF
Ask AI
Cite

Copy citation

DOINot available for this paper

Abstract

Group attention, introduced by the time series forecasting model Chronos-2, attends over the variates of a group at a fixed patch index and serves both multivariate (MV) and in-context learning (ICL) forecasting. Rather than evaluating this cross-variate attention design as a whole, we ask which part of the mechanism earns the benefit and probe its applicability to both MV and ICL regimes. By editing the attention matrix αα at inference we separate the two pathways a head comprises: V/O, which projects a weighted summary of the group, and Q/K, which decides the weights. Uniform pooling (V/O without any Q/K weighting) is positive on 18 of our 20 sensor-network configurations, while the learned weighting (Q/K) splits by group type: its contribution is positive or negligible for MV, but materially degrades 8 of the 10 sensor-network ICL configurations, leaving 4 of them worse than univariate inference. By isolating the impact of different layers, we find that uniforming αα in the first block alone improves every ICL configuration we test.

Figures & tables

Figure 1: The four α\alpha settings of § 2 : darker cells are larger weights, and consecutive differences isolate pooling (V/O) from routing (Q/K).
Figure 2: Each contribution on five sensor networks, each as an MV and an ICL group at two protocols. Bars are 95% paired bootstrap intervals; filled markers are material (§ 2 ).
Figure 3: Where the learned weighting’s cost lives, at the ctx-204 protocol. (a) one block active, α=I\alpha=I elsewhere, and (c) one block forced uniform against the released Chronos-2 — both plotting independent single-block edits as unconnected markers. (b) the leading blocks active cumulatively, its nested prefixes connected; (a) and (b) run against univariate inference and share a yy axis.
Appendix figures & tables6 assets

Supplementary material from the paper’s appendix.

Appendix

DatasetGroupnnctxm/Nm/NPoolingShuffledLearnedTotal
Sensor networks
Atlantic tidesMV2732040.72+5.16∗+5.16^{*}+1.15∗+1.15^{*}+8.15∗+8.15^{*}+14.47∗+14.47^{*}
Atlantic tidesICL3002040.41+7.59∗+7.59^{*}+2.00∗+2.00^{*}−1.85∗-1.85^{*}+7.74∗+7.74^{*}
Atlantic tidesMV1683200.72+5.41∗+5.41^{*}+1.19∗+1.19^{*}+8.37∗+8.37^{*}+14.97∗+14.97^{*}
Atlantic tidesICL3003200.41+10.89∗+10.89^{*}+2.72∗+2.72^{*}−3.81∗-3.81^{*}+9.80∗+9.80^{*}
Pacific tidesMV2002040.55+5.26∗+5.26^{*}+2.82∗+2.82^{*}+2.31∗+2.31\phantom{{}^{*}}+10.40∗+10.40^{*}
Appendix
Table A1: All 39 configurations. Contributions in points of univariate-baseline MAE; the three terms add to the total. ∗ marks a material effect: a 95% paired bootstrap interval excluding ±0.1\pm 0.1 points. nn is the number of forecast windows; ctx is context length. The sensor-network rows are the data plotted in Figure 2 .
   error vs univariate (points)
ConfigurationGroupnnraw meanmatcheduniformlearned
tidesp256_iclICL250−216.8-216.8−14.1-14.1+7.71+7.71+7.17+7.17
netpjmp256_iclICL200−98.3-98.3−8.7-8.7−0.87-0.87−3.77-3.77
tidesp256_mvMV250−222.1-222.1−51.0-51.0+5.18+5.18+14.59+14.59
netuscrnp256_mvMV200−63.6-63.6−9.3-9.3+1.79+1.79+3.35+3.35
Appendix
Table A2: Is the V/O projection necessary? Each column is that condition’s error relative to the univariate baseline, in points. raw mean substitutes the unprojected cross-variate mean for WOVW_{OV} ; matched rescales that mean to the norm of the learned update; uniform and learned are the settings of § 2 . Replacing the projection costs far more than the entire effect the paper reports, so the uniform- α\alpha result is not a statement that cross-variate attention reduces to averaging.
  one blockleading blockslate
DatasetGroup037110–10–30–50–7all8–11
Atlantic tidesICL+4.75∗+4.75^{*}+0.06∗+0.06\phantom{{}^{*}}+0.12∗+0.12\phantom{{}^{*}}+0.07∗+0.07\phantom{{}^{*}}+4.49∗+4.49^{*}+0.70∗+0.70\phantom{{}^{*}}−0.61∗-0.61\phantom{{}^{*}}+1.74∗+1.74^{*}−1.84∗-1.84^{*}+0.15∗+0.15\phantom{{}^{*}}
Pacific tidesICL+1.18∗+1.18\phantom{{}^{*}}+0.15∗+0.15^{*}+0.52∗+0.52^{*}+0.22∗+0.22^{*}−1.34∗-1.34\phantom{{}^{*}}−8.97∗-8.97^{*}−9.77∗-9.77^{*}−8.05∗-8.05^{*}−16.09∗-16.09^{*}+0.88∗+0.88^{*}
USCRNICL+0.34∗+0.34\phantom{{}^{*}}−0.01∗-0.01\phantom{{}^{*}}−0.01∗-0.01\phantom{{}^{*}}+0.08∗+0.08\phantom{{}^{*}}−1.78∗-1.78^{*}−2.26∗-2.26^{*}−3.41∗-3.41^{*}−0.86∗-0.86\phantom{{}^{*}}+0.14∗+0.14\phantom{{}^{*}}+0.10∗+0.10\phantom{{}^{*}}
PJM loadICL−1.82∗-1.82^{*}+0.00∗+0.00\phantom{{}^{*}}+0.01∗+0.01\phantom{{}^{*}}+0.04∗+0.04\phantom{{}^{*}}−4.12∗-4.12^{*}−6.15∗-6.15^{*}−6.71∗-6.71^{*}−3.55∗-3.55^{*}−3.34∗-3.34^{*}+0.02∗+0.02\phantom{{}^{*}}
EU loadICL−2.23∗-2.23^{*}+0.06∗+0.06\phantom{{}^{*}}+0.13∗+0.13\phantom{{}^{*}}+0.11∗+0.11\phantom{{}^{*}}−6.67∗-6.67^{*}−9.08∗-9.08^{*}−7.70∗-7.70^{*}−4.14∗-4.14^{*}−8.20∗-8.20^{*}+0.68∗+0.68^{*}
Appendix
Table A3: The restriction sweep: learned weighting when group attention is active only in the named blocks and α=I\alpha=I in the rest, against univariate inference, in points of univariate-baseline MAE. ∗ marks a material effect (§ 2 ). Left group: one block in isolation. Middle: the leading blocks together. Right: the last four blocks only. In isolation no block other than 0 moves the forecast by more than 0.520.52 points either way, while block 0 spans −2.23-2.23 to +6.44+6.44 . Of the 5 multi-block subsets on 5 in-context datasets, 18 of 25 are materially negative and the worst reaches −16.09-16.09 ; on multivariate datasets 1 of 25 are, and the worst is −2.43-2.43 . With only the last four blocks active, no dataset moves by more than 0.880.88 points.
DatasetGroupblock 0 onlyall twelve
Atlantic tidesICL+0.20∗+0.20\phantom{{}^{*}}+7.59∗+7.59^{*}
Pacific tidesICL+1.14∗+1.14^{*}+13.77∗+13.77^{*}
USCRNICL−0.04∗-0.04\phantom{{}^{*}}+1.71∗+1.71^{*}
PJM loadICL−0.06∗-0.06\phantom{{}^{*}}−0.87∗-0.87\phantom{{}^{*}}
EU loadICL+0.10∗+0.10\phantom{{}^{*}}−0.05∗-0.05\phantom{{}^{*}}
Atlantic tidesMV+0.21∗+0.21^{*}+5.16∗+5.16^{*}
Appendix
Table A4: The other half of the restriction sweep: pooling (the V/O pathway) with one block active against all twelve, in points of univariate-baseline MAE. ∗ marks a material effect. Pooling grows with the number of active blocks on 8 of 10 datasets, the opposite of the learned weighting in Table A3 .
     block 0 uniformall blocks
DatasetGroupctxnnUV MAEMAEWQLcoverageMAE
In-context groups
Atlantic tidesICL2043000.14040.1404+6.27∗+6.27^{*}+5.56∗+5.56^{*}0.751→0.7580.751\to 0.758−0.15∗-0.15\phantom{{}^{*}}
Atlantic tidesICL3203000.17050.1705+6.45∗+6.45^{*}+6.58∗+6.58^{*}0.700→0.7760.700\to 0.776+1.09∗+1.09\phantom{{}^{*}}
Pacific tidesICL2042000.12280.1228+20.49∗+20.49^{*}+19.53∗+19.53^{*}0.677→0.8270.677\to 0.827+11.90∗+11.90^{*}
Pacific tidesICL3202000.14860.1486+10.85∗+10.85^{*}+13.38∗+13.38^{*}0.593→0.7080.593\to 0.708+10.46∗+10.46^{*}
Appendix
Table A5: The repair intervention on every sensor-network configuration: α\alpha forced uniform in the named blocks of the otherwise-shipped model, against the shipped model itself. Positive is better than shipped, in points of univariate-baseline MAE (or of the univariate WQL). ∗ marks a material effect: a 95% paired bootstrap interval excluding ±0.1\pm 0.1 points, the floor calibrated by the block-5 near-null. Coverage is the empirical coverage of the nominal 80%80\% interval, shipped →\to repaired. Repairing block 0 alone improves every in-context configuration on all three metrics; its only material costs are on tidal multivariate groups, which are also the only configurations where learned weighting materially helps (Table A1 ).
   blocks forced uniform
DatasetGroupctx01357110–10–3all
Atlantic tidesICL204+6.27∗+6.27^{*}+3.57∗+3.57^{*}−0.00∗-0.00\phantom{{}^{*}}−0.01∗-0.01\phantom{{}^{*}}−0.94∗-0.94^{*}−1.46∗-1.46^{*}+1.57∗+1.57^{*}+0.62∗+0.62\phantom{{}^{*}}−0.15∗-0.15\phantom{{}^{*}}
Atlantic tidesICL320+6.45∗+6.45^{*}+5.72∗+5.72^{*}+0.02∗+0.02\phantom{{}^{*}}−0.00∗-0.00\phantom{{}^{*}}−0.31∗-0.31\phantom{{}^{*}}−0.64∗-0.64^{*}+2.89∗+2.89^{*}+1.95∗+1.95^{*}+1.09∗+1.09\phantom{{}^{*}}
Pacific tidesICL204+20.49∗+20.49^{*}+1.80∗+1.80^{*}−0.08∗-0.08\phantom{{}^{*}}−0.03∗-0.03\phantom{{}^{*}}−3.10∗-3.10^{*}−3.90∗-3.90^{*}+15.57∗+15.57^{*}+14.26∗+14.26^{*}+11.90∗+11.90^{*}
Pacific tidesICL320+10.85∗+10.85^{*}+3.00∗+3.00^{*}+0.30∗+0.30^{*}−0.04∗-0.04\phantom{{}^{*}}−0.54∗-0.54^{*}−1.61∗-1.61^{*}+12.76∗+12.76^{*}+11.87∗+11.87^{*}+10.46∗+10.46^{*}
USCRNICL204+0.58∗+0.58\phantom{{}^{*}}+0.92∗+0.92^{*}−0.50∗-0.50\phantom{{}^{*}}−0.03∗-0.03\phantom{{}^{*}}+0.39∗+0.39^{*}−0.27∗-0.27^{*}−0.41∗-0.41\phantom{{}^{*}}−0.49∗-0.49\phantom{{}^{*}}−0.44∗-0.44\phantom{{}^{*}}
Appendix
Table A6: Which block has to be neutralized, on all 20 configurations. Each column forces α\alpha uniform in the named block or blocks of the otherwise-shipped model; positive is better than shipped, in points of univariate-baseline MAE. ∗ marks a 95% paired bootstrap interval excluding ±0.1\pm 0.1 points, the floor this column calibrates. Nearly all of the recoverable effect sits in blocks 0 and 1, and no single later block recovers anything. Block 5 is the intended null: its attention is already near-uniform in the shipped model, so forcing it uniform moves the forecast by at most 0.0810.081 points, which is material on none of the 20 — though 14 of the 20 intervals do exclude zero, which is what calibrates the floor and why effects are reported against it rather than against zero.

Explore similar work

CardsList

  1. When and Why Grouping Attention Heads Accelerates Muon Optimization

    May 9, 2026Hongtao Zhang, Wenjie Zhou, Wei Chen +1Multi-Head AttentionMuon

    State Key Laboratory of AI Safety, Institute of Computing Technology, Chinese Academy of Sciences · School of Advanced Interdisciplinary Sciences, University of Chinese Academy of Sciences · University of Chinese Academy of Sciences

  2. In-context Learning of Single-index Targets: Comparing Kernel and Feature Learners

    Oct 1, 2026Haotian Gu, Yizhou Xu, Lenka ZdeborováIn-Context LearningLinear Attention

    Statistical Physics of Computation Laboratory, École Polytechnique Fédérale de Lausanne (EPFL) · University of Chinese Academy of Sciences · Information, Learning and Physics Laboratory, École Polytechnique Fédérale de Lausanne (EPFL)

  3. GAttNHP: Group Attention Neural Hawkes Process for Extrapolation Reasoning in Temporal Knowledge Graphs

    Jul 16, 2026Xiangni Tian, Kaixian Yu, Runpeng Dai +2Temporal Knowledge GraphsLong-Range Temporal Dependencies

    Yunnan Key Laboratory of Statistical Modeling and Data Analysis, Yunnan University, Kunming, China · Insilicom LLC, Peabody, MA, USA · University of North Carolina at Chapel Hill, North Carolina, USA