PaperlayerView exact revision on arXiv

Source-linked AI summary

Attention Function as an Intrinsic Inductive Bias: How Models' Behavior Diverges in Novel Contexts

Dong Gyun Kang, Megha Thukral, Kwangsoo Kim

arXiv:2609.39188v1cs.LG

TL;DR

The paper asks whether attention activation can act as a given inductive bias in Transformers, analogous to priors that precede experience. It proposes MoFA, which fixes a mixture of softmax and sigmoid heads before training, and finds that the ratio’s influence is muted in-distribution but re-emerges under OOD shift along a domain-structure axis.

  • Problem

    The paper asks whether attention activation can serve as an intrinsic inductive bias in Transformers and shape generalization to novel contexts.

    Source: §1 · p. 1

  • Method

    MoFA is a parameter-free modification that fixes a ratio of softmax and sigmoid attention heads before training.

    Source: §1 · p. 2, §3.2 · p. 3

  • Results

    Under distribution shift, activation-ratio effects re-emerge unevenly, with 78.3% of domain response tracking a single axis of document-level relevance structure.

    Source: §5 · p. 7, §6 · p. 8

  • Takeaways & Limitations

    The results support treating activation choice as an intrinsic prior whose influence is masked in-distribution and re-emerges under distribution shift.

    Source: §6 · p. 8

    Experiments use one 124M-parameter model, one pretraining corpus, and one activation pair, leaving generalization to other scales, distributions, and activation combinations open.

    Source: §5 · p. 8

Abstract

from arXiv · show

Developmental psychology holds that certain priors are given to infants prior to experience rather than induced from data, and that the influence of such priors is suppressed under strong, well-constrained conditions but reasserts itself under weak ones. We ask whether an analogous principle holds for the Transformer: can the activation function given to attention heads serve as an intrinsic inductive bias? We propose Mixture of Function Attention (MoFA), a parameter-free modification to multi-head attention that fixes a ratio of softmax and sigmoid heads before training. Across five ratios, a 124M-parameter GPT-2 model, and five seeds, we find that this given ratio has little effect in-distribution -- differences between ratios are statistically negligible for moderate mixtures and remain small even at the extremes -- but its influence re-emerges sharply under zero-shot distribution shift across 15 out-of-distribution domains. Perplexity gaps between ratios widen by more than an order of magnitude on several domains, and the best-performing ratio tracks a single axis of domain structure, separating short, informal text (softmax-favoring) from technical, long-form text (sigmoid-favoring), that explains 78.3% of the variance in domain response. This reorganization is visible at the head level: sigmoid heads show an accelerating drop in attention entropy as their ratio increases, while softmax heads respond more modestly, yielding a consistent division of labor between the two head types. Our results suggest that activation choice functions as a given prior whose influence is masked in-distribution and re-emerges out-of-distribution.

1 Introduction

Inspired by accounts of prior knowledge in developmental psychology, the paper asks whether attention activation can act as a given inductive bias in Transformers. It proposes MoFA and reports that the influence of the fixed activation ratio is muted in-distribution but re-emerges under distribution shift.

1 Introduction

  • Developmental psychology describes certain knowledge systems as present before experience, motivating the question of whether attention activation can serve as a Transformer’s intrinsic inductive bias.The paper frames the analogy around how a prior may shape representation and generalization to novel contexts.Source: §1 · p. 1
  • The authors hypothesize that well-constrained in-distribution training may mask differences between attention ratios, while zero-shot OOD settings may reveal them.The proposed contrast is between task-specific training constraints and their absence on a new distribution.Source: §1 · p. 1
  • MoFA mixes attention activations across heads at a fixed ratio, and the paper reports that this given prior is suppressed in-distribution but re-emerges under distribution shifts.The paper also reports that no single activation ratio is universally optimal across 15 OOD domains and that head types adopt complementary roles.Source: §1 · p. 2

2 Related Work

Related work examines architectural priors, alternatives to softmax, head-level heterogeneity, and kernel interpretations of attention. The paper distinguishes MoFA by varying activation functions across heads within dot-product attention.

Developmentally Inspired Architectures

  • Developmentally inspired models encode priors such as lattice symmetry or spatial locality as hard architectural constraints tied to specific domains.The paper relates its premise to these approaches while studying a mixture of attention activations.Source: §2 · p. 2

Modifying Softmax Attention

  • Other approaches to softmax attention used kernel approximations, sigmoid substitution, or head-specific gating, but applied the choice uniformly across heads.MoFA instead makes activation a per-head decision.Source: §2 · p. 2

Head-Level Heterogeneity

  • Prior methods varied attention mechanisms or routed among heads, but did not introduce activation-function diversity within the same dot-product attention structure.The cited examples include hybrid full/linear attention and dynamic head routing.Source: §2 · p. 2

Attention Head as Kernel

  • Kernel views cast attention heads as kernel smoothers or learning procedures, motivating activation choice as a way to select the function space a head searches.The cited work relates attention to kernel smoothing, reproducing kernels, and functional gradient descent.Source: §2 · p. 3

3 Method

Standard multi-head attention applies row-wise softmax to scaled query–key scores and uses the resulting weights over values. MoFA assigns softmax or sigmoid scoring activations to different heads at a fixed ratio before training.

3.1 Preliminaries

  • Standard scaled dot-product attention uses queries, keys, and values; row-wise softmax produces a distribution over positions, and multiple heads operate on independent projections.Head outputs are concatenated and projected back.Source: §3.1 · p. 3

3.2 Mixture of Function Attention

  • MoFA partitions heads into softmax and sigmoid groups at a fixed ratio, with the same contiguous assignment in every layer before training.Softmax heads allocate attention competitively with rows summing to 1, while sigmoid heads gate positions independently for dense, distributed patterns; the heads share QKV projections and add no parameters.Source: §3.2 · p. 3

4 Experiments

Across controlled ratio configurations, in-domain differences were small, while zero-shot performance varied across OOD domains. The experiments also trace this variation to domain structure and distinct entropy responses in softmax and sigmoid heads.

4.1 Setup

  • Across five ratios, the 124M-parameter GPT-2 models were compared with identical data order, initialization, and hyperparameters across five seeds.Models were pretrained on OpenWebText and evaluated across 15 OOD domains without fine-tuning or domain adaptation.Source: §4.1 · p. 4, §4.3 · p. 4

4.2 In-Domain Validation.

  • On OpenWebText, Hsig=3 matched the baseline at 21.26 validation PPL (p = 0.693), while Hsig=12 degraded by 0.53 PPL (21.26 → 21.79).The maximum degradation was statistically significant for Hsig ≥6 (p < 0.01).Source: §4.2 · p. 4, Table 1 · p. 4

4.3.1 Performance by Domain.

  • Zero-shot perplexity varied by domain: each ratio was best on at least one domain, and softmax tended to suit short, informal text while sigmoid suited technical, long-form text.Conversational domains fell between these preferences.Source: §4.3.1 · p. 5, Table 2 · p. 5
  • Across 15 OOD domains, no ratio dominates; pure endpoints were worst in 13/15 domains, and their worst-place gaps were significant in 11/13 cases.By contrast, when a mixed ratio was worst, its gap from the best configuration was not statistically significant.Source: §4.3.1 · p. 5, Table 2 · p. 5
  • A single domain-response axis explained 78.3% of variance, separating short units with concentrated relevance from long documents where relevance spans many positions.LAMBADA and Empathetic depart from this account, including cases where surface form and long-range dependency pull in different directions.Source: §4.3.1 · pp. 5–6, Figure 1 · p. 6

4.4.1 Consecutive-Ratio Comparisons

  • Sigmoid-head entropy decreased significantly at every adjacent-ratio step, with the decline accelerating to −0.36 from H3:9 to H0:12 (p < 10−4).Softmax-head entropy changed significantly only from H6:6 to H3:9 (p = 0.035); the first two transitions were not significant.Source: §4.4.1 · p. 6, Table 3 · p. 7

4.4.2 Where This Occurs in the Network

  • Across layers, softmax-head entropy stayed mostly stable, while sigmoid-head entropy dropped sharply only at H0:12, creating a deeper mid-network trough.The figure averages attention entropy across 15 OOD domains and separates softmax heads (left) from sigmoid heads (right).Source: §4.4.2 · p. 7, Figure 2 · p. 7

5 Discussion

Under distribution shift, activation-ratio preferences align with a single axis of domain structure rather than varying idiosyncratically. Head-level and domain-level asymmetries coincide, although the authors lack a mechanistic explanation and note important experimental scope limits.

5 Discussion

  • 78.3% of domain response collapses onto one axis tracking local versus distributed document-level relevance structure.LAMBADA and Empathetic depart from this dominant dimension, so the axis is not exhaustive.Source: §5 · p. 7
  • Pure endpoints are the worst-performing configuration in 13/15 domains, while sigmoid heads reorganize continuously and softmax heads remain stable until they cease being the majority.The authors suggest head interactions may underlie the shared asymmetry, but say they lack a mechanistic explanation.Source: §5 · pp. 7–8
  • The experiments cover one 124M-parameter model, one pretraining corpus, and softmax–sigmoid mixtures; layer-level analysis is descriptive.Whether the pattern holds at larger scales, on other pretraining distributions, or with other activation pairs remains open.Source: §5 · p. 8

6 Conclusion

The paper finds that attention activation choice has little influence in-distribution but reasserts itself under distribution shift, tracking domain structure and reorganizing attention at the head level.

6 Conclusion

  • Under distribution shift, activation choice reasserts itself sharply and unevenly, tracking one axis of domain structure and reorganizing attention at the head level.The authors frame activation choice as an intrinsic prior whose influence is largely suppressed in-distribution.Source: §6 · p. 8

7 Broader Impact

The work studies how mixtures of attention activation functions across heads affect language-model perplexity under distribution shift, using core-knowledge accounts of inductive bias to motivate the design and evaluation.

7 Broader Impact

  • The study evaluates how mixing activation functions across attention heads affects language-model perplexity under distribution shift, motivated by core-knowledge accounts of inductive bias.Source: §7 · p. 8

8 Reproducibility

The experiments use shared configurations and fixed seeds to support reproducibility, with public code and data and a planned release of code, checkpoints, and evaluation scripts.

8 Reproducibility

  • All five configurations share data order, initialization, and hyperparameters across five fixed seeds, differing only in sigmoid-head count; code and data are public, with artifacts planned for release upon acceptance.The 124M GPT-2 setup, training procedure, and optimizer settings are specified, and OOD datasets are used only for zero-shot perplexity measurement.Source: §8 · p. 8

A Dataset Details

Table 4 catalogs the dataset identifiers, splits, and source publications for the 15 out-of-distribution domains, with evaluation performed zero-shot.

A Dataset Details

  • Table 4 maps the 15 out-of-distribution domains to dataset identifiers, splits, and original source publications.All domains are evaluated zero-shot, using the custom 50k-vocabulary BPE tokenizer and block chunking strategy.Source: Table 4 · p. 11, A Dataset Details · p. 11
  • Novel and Reddit use train[:5%] splits, while other domains use canonical test or validation splits.SST-2 uses validation; the other domains use their canonical test splits.Source: A Dataset Details · p. 11
  • arXiv and PubMed share a summarization dataset construction but draw from disjoint document collections.The collections are arXiv preprints and PubMed abstracts, respectively.Source: A Dataset Details · p. 11

Related papers

Loading 2609.39188v1…