Pith. sign in

REVIEW 2 major objections 5 minor 25 references

Small transformer architectures for task switching

T0 review · 2 major / 5 minor · reviewed 2026-08-05 · deepseek-v4-flash

Pith's one-line read A position-aware transformer with rational attention hits 95% on a task-switching benchmark that stumps standard transformers, LSTMs, and MLPs.

desk verdict A new toy benchmark and a potentially interesting expressive-attention result, but the headline comparison is confounded by a 5-15x parameter gap; needs matched-size retraining before the claims are credible. read the letter →

arxiv 2508.04461 v1 pith:UDTHGP2Q submitted 2025-08-06 cs.LG cs.AI

classification cs.LGcs.AI
keywords taskswitchingtransformerattentioncisformerexpressiveIARCbenchmarksmallmodelsin-contextlearning
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

The paper asks whether attention-based architectures can beat classical models in a small-scale setting, where scaling laws do not decide the outcome. To answer it, the author introduces the IARC task-switching benchmark: an ongoing stream of encoded numbers interrupted by control tokens that command increment, addition, reverse copy, and context-dependent steps. Across model classes, standard transformers, LSTMs, and MLPs all plateau at modest accuracies, with the standard transformer at 45–58% depending on the attention formula. The combination of the cisformer — a transformer whose weights are not shared along the context dimension — with expressive attention, a rational square of the query-key dot product in place of softmax, reaches around 95% accuracy. Why this matters: it shows that in this regime the functional form of attention, not just model scale or architecture class, can determine whether an algorithmically structured task is learnable.

What carries the argument

The cisformer, a transformer in which all parameter tensors are position-dependent instead of being broadcast across context positions; the author motivates it by noting that control tokens arrive stochastically, so past context slots carry different information. Expressive attention (EA), defined by the rational attention weight $A(z_{ij}) \sim z_{ij}^2/(1+z_{ij}^2)$ instead of the usual softmax $A(z_{ij}) \sim \exp(\beta z_{ij})$, where $z_{ij}=Q_i\cdot K_j$. The paper argues EA makes attention weights small when queries and keys are orthogonal rather than only when anti-parallel, which enriches the geometry available to attention heads. Together, EA supplies a biquadratic dependence on to

What would settle it

Train a standard transformer raised to the cisformer's roughly 1.3M parameters on the same IARC generator; if it approaches the 95% accuracy of cisformer + expressive attention, the paper's claim that standard transformers cannot solve IARC would be falsified. A second check: run the cisformer with softmax attention; if it too reaches around 95%, the expressive-attention component would not carry the result.

Watch

Extended reading notes

Core claim

The central empirical claim is that the IARC task is solvable at small scale only when attention abandons translational invariance and softmax. With 60 layers, a standard transformer using dot-product attention reaches 45% on the combined task; its LSTM and MLP counterparts score similarly. Swapping softmax for expressive attention lifts the standard transformer only to 58%. Replacing the shared-position weights with the cisformer while keeping softmax leaves it in the same modest range as the other classical models. Only the cisformer with expressive attention reaches around 95%, making this combination the paper's central discovery.

Load-bearing premise

The conclusion that standard transformers 'cannot solve' IARC assumes the 240K-parameter standard transformer is a fair comparison against the 1.3M-parameter cisformer, an assumption the paper states in Section 6 as 'overall model size was comparable'.

Editorial extensions

If this is right

  • On the IARC benchmark, task-switching ability is not a generic property of attention: a standard transformer with softmax reaches only 45% and with expressive attention 58%.
  • Per-position parameter matrices (the cisformer) plus expressive attention reach about 95%, so breaking translational invariance is sufficient to solve IARC at small scale.
  • LSTM and MLP baselines with larger parameter counts (up to 3.7M) also plateau, so the gap is not simply a matter of giving the winner more parameters within the tested range.
  • Changing softmax to the rational quadratic expression changes an architecture from failing to succeeding, both for the standard transformer and for the cisformer.
  • IARC provides a compact, configurable evaluation protocol for testing attention variants on task switching without large-scale training.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • A parameter-matched control — widening or deepening the standard transformer to the cisformer's 1.3M parameters — would determine whether the gap is due to position-dependent weights or to raw capacity; the paper does not report that run.
  • Because IARC uses a 6±3 spacing between control tokens and a context length of 24, the cisformer's per-position advantage may be specific to short, irregular-but-bounded contexts; testing longer or stochastic horizons would delimit the result.
  • The rational weighting $z^2/(1+z^2)$ could plausibly be ported to linear-attention variants that already avoid softmax, potentially yielding a sub-quadratic model with some of EA's task-switching benefit; this is an extension the paper does not claim.
  • The 'control tape' input representation places control tokens in a separate embedding subspace; a natural follow-up is to check whether interspersing control tokens directly in the symbol stream preserves the 95% result.
Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

2 major / 5 minor

Summary. The paper introduces IARC, a small-scale task-switching benchmark in which an ongoing token sequence is governed by stochastically appearing control tokens for increment/addition/reverse-copy/context subtasks. The author compares standard transformers, a position-wise variant called the cisformer, an MLP, and an LSTM, with either standard dot-product attention (DPA) or the recently proposed expressive attention (EA). The central empirical claim is that standard transformers cannot solve IARC, that the other classical models also achieve only modest accuracy, and that the cisformer with EA is the only tested configuration reaching substantial performance (~95%). The paper also argues that this points to attention-space geometry as an important design axis.

Significance. If the comparison were fair and the results held, the paper would be a useful small-scale probe of when attention-based architectures beat MLPs and RNNs, and it would provide evidence that non-translation-invariant, per-position parameterization plus a rational attention kernel enables a task-switching capability that standard transformers lack. The IARC benchmark is clearly specified, the ablations over IAR/IA/IR are informative, and the author reports parameter counts for each model, which is commendable. However, the central comparison is currently undermined by a large parameter imbalance between the standard transformer and all other models, and the author's own §4.1 explicitly raises limited parameter count as a possible explanation for the standard transformer's poor performance without testing it. The contribution is therefore not yet established, though the issue is empirical and could in principle be addressed with additional experiments.

major comments (2)
  1. [§5 and §6] The fair-comparison claim is contradicted by the reported parameter counts. Section 5 states that the standard transformer has 60 layers with d=20, i.e., about 240K parameters, while the cisformer has 1.3M parameters, the MLP 1.9M, and the LSTM 3.7M. Section 6 says layers were chosen so that 'overall model size was comparable,' but the listed choice (60 transformer layers vs 12/16 cisformer/MLP layers and 2 LSTM layers) produces a factor of 5–15 difference, not comparability. This matters because §4.1 explicitly suggests that the standard transformer's poor IARC performance 'could be the limited number of available adaptable parameters, namely 60·4000=240K.' The abstract's claim that standard transformers cannot solve IARC is therefore unsupported: the standard transformer is given a much smaller capacity budget than every model it is compared with. A proper control would train a standar
  2. [Table 1 and Fig. 1] The headline comparison changes architecture and parameter budget simultaneously. The 95% result is for cisformer+EA at 1.3M parameters, whereas the standard transformer with EA in Table 1 reaches only 0.58 on IARC at 240K parameters. Thus, the boost cannot be attributed to expressive attention alone; it may partly reflect the much larger parameter budget of the cisformer. The comparison between cisformer+DPA and cisformer+EA at the same 1.3M size does isolate the attention mechanism, and that is a meaningful result, but the broader claim that 'a combination of the latter [cisformer and expressive attention] is the only model able to achieve considerable performance' is overreaching. The paper should present a factorial design (standard vs cisformer × DPA vs EA) at matched parameter counts, or clearly restrict the conclusion to the specific configurations tested.
minor comments (5)
  1. [Abstract and §6] The abstract refers to 'extensive attention' and §6 refers to the 'cistransformer'; both should be 'expressive attention' and 'cisformer'.
  2. [§4.1] The per-layer parameter estimate is given as Nlayer = 11d^2 + O(d), which for d=20 is about 4400, not 4000. The text says 'about 4000' and then derives 60·4000=240K. Please reconcile the formula and the numerical value.
  3. [§5] The paper states that results are independent of initialization and seed, but no multiple-seed statistics or error bars are reported. Given the visible fluctuations in Fig. 1, a short statement of the number of seeds and variance would strengthen the empirical claims.
  4. [§3] The IARC generation protocol is described for a 6±3 control-token interval distribution, but it is not stated whether this is applied separately to each control-token type or whether the four control tokens are first sampled uniformly and then placed with this spacing. Clarifying this would make the benchmark fully reproducible.
  5. [General] No code or data availability statement is included. Since the contribution includes a new benchmark, making the sequence-generation and training code available would substantially increase its usefulness.

Circularity Check

0 steps flagged · score 0.0 of 10

No circularity found; the paper's central claim is an empirical comparison, not a derivation that reduces to its own inputs.

full rationale

The paper does not present a formal derivation of its main result; it reports empirical measurements of different architectures on a new task-switching benchmark (IARC). The winning configuration (cisformer + expressive attention) is defined via Eq. (8) and its performance is measured in Fig. 1 and Table 1, not derived from the benchmark definition or from the cited prior work. The self-citation to Gros (2024) motivates expressive attention and provides a geometric intuition, but the decisive evidence is the new experimental evaluation carried out here, so the citation is not load-bearing in a circular sense. The IARC benchmark is introduced by the author, but that fact alone does not make the empirical comparison circular: the benchmark has an independent input-output rule (Sect. 3) and the architectures are evaluated against that rule. A legitimate methodological concern is raised by the parameter-count disparity (standard transformer 240K versus cisformer 1.3M, MLP 1.9M, LSTM 3.7M) and by Section 6's claim that 'overall model size was comparable' despite these counts differing by factors of 5–15; however, this is a fairness/correctness issue, not a circularity where a prediction is equivalent to a fitted input by construction. No equation in the paper reduces the claimed result to its own assumptions, no fitted parameter is renamed as a prediction, and no uniqueness theorem is imported from the authors' prior work to forbid alternatives. Therefore the paper receives a circularity score of 0.

Assumptions & free parameters 5 free parameters · 5 assumptions · 2 invented entities

The central comparison rests on several hand-designed choices: the task definition, the control token distribution, the embedding scheme, and the training hyperparameters. Most importantly, the model size comparison is not actually fair, as parameter counts vary 5-15x across architectures. The cisformer and expressive attention are new architectural entities introduced to improve performance, and their only evidence comes from this same paper.

free parameters (5)
  • Learning rate and momentum = 0.02, 0.8
    Chosen by hand and used for all models; they affect convergence and final performance.
  • Number of layers per model = Transformer 60, cisformer 12, MLP 16, LSTM 2
    Selected to approximate equal model size, but actually results in parameter counts differing by 5-15x, which directly affects the comparison.
  • Context length Ncon = 24
    Chosen to accommodate the maximum distance between control tokens (2*9=18), affecting memory and task difficulty.
  • Batch size and epochs = 200, 8000
    Fixed by hand for all runs; these training choices influence final accuracy.
  • Control token interval distribution = 6 +/- 3, equal probabilities
    Task design choice that sets how often tasks switch; changes in this distribution would alter the difficulty.
assumptions (5)
  • domain assumption IARC task definition (Eq. 1-4): I/A/R/C operations as specified
    The benchmark's structure is assumed to be a meaningful test of task switching; no justification is given for why these specific operations are representative.
  • domain assumption Control tokens are stochastically interspersed with 6 +/- 3 spacing and equal probabilities
    This specific stochastic process is chosen by the author and is not derived from any principle; it determines the task's difficulty.
  • domain assumption One-hot embedding with control tape (Eq. 5) and d = N + Ncontrol
    The encoding scheme is a design choice; other embeddings could change relative model performance.
  • domain assumption ALiBi positional encoding (Press et al., 2021) is used for all attention models
    Positional encoding choice can affect how models handle context; the paper does not ablate this choice.
  • domain assumption Training with cross-entropy loss and SGD with momentum
    The paper states 'learning rate and momentum' but does not explicitly name the optimizer; this is a standard assumption but not stated precisely.
invented entities (2)
  • Cisformer
    purpose: A transformer variant where parameters are not shared across context positions, giving each position its own weights.
    This architecture is introduced in the paper as a way to enlarge model size while keeping computation structure similar. It has no external falsifiable prediction beyond the benchmark results shown.
  • Expressive attention (EA)
    purpose: An alternative attention weight function A(z) = z^2/(1+z^2) replacing softmax.
    EA is proposed in the author's prior work (Gros 2024). In this paper it is applied to IARC, but there is no independent handle outside the author's own benchmarks. The paper claims it improves expressivity but provides no external test.

how reviews work

0 comments
Cite this review

Pith. "Pith review of Small transformer architectures for task switching." pith.science (2026). https://pith.science/paper/UDTHGP2Q

@misc{pith2026250804461,
  author       = {Pith},
  title        = {Pith review of: Small transformer architectures for task switching},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/UDTHGP2Q}},
  note         = {Machine review of arXiv:2508.04461}
}
read the original abstract

The rapid progress seen in terms of large-scale generative AI is largely based on the attention mechanism. It is conversely non-trivial to conceive small-scale applications for which attention-based architectures outperform traditional approaches, such as multi-layer perceptrons or recurrent networks. We examine this problem in the context of 'task switching'. In this framework models work on ongoing token sequences with the current task being determined by stochastically interspersed control tokens. We show that standard transformers cannot solve a basic task switching reference model based on finite domain arithmetics which contains subtasks dedicated to increment / addition / reverse copy / context (IARC). We show that transformers, long short-term memory recurrent networks (LSTM), and plain multi-layer perceptrons (MLPs) achieve similar, but only modest prediction accuracies. We enlarge our comparative study by including an extension of the standard transformer architecture to its non-translational invariant counterpart, the cisformer, and an alternative attention mechanism, extensive attention. A combination of the latter is found to be the only model able to achieve considerable performance levels, of around 95%. Our results indicate that the workings of attention can be understood better, and even improved, when comparing qualitatively different formulations in task-switching settings.

Figures

Figures reproduced from arXiv: 2508.04461 by the authors.

Figure 1
Figure 1. Evaluation of the IARC task switching framework defined in Sect. 3. Shown are results for a LSTM recurrent network (black), a MLP (green), and two versions of cisformers (see Sect. 4.2), namely with standard attention (blue), and with expressive attention (red). The equivalent results for classical transformers are listed in [PITH_FULL_IMAGE:figures/full_fig_p002_1.png] view at source ↗

Discussion (0). Sign in to comment.

Reference graph

Works this paper leans on

25 extracted references · 13 canonical work pages

  1. [1]

    Advances in Neural Information Processing Systems 36 (2024)

    Bietti, A., Cabannes, V., Bouchacourt, D., Jegou, H., Bottou, L.: Birth of a transformer: A memory viewpoint. Advances in Neural Information Processing Systems 36 (2024)

  2. [2]

    ACM Computing Surveys56(12), 1–32 (2024)

    Chen, S., Zhang, Y., Yang, Q.: Multi-task learning in natural language process- ing: An overview. ACM Computing Surveys56(12), 1–32 (2024)

  3. [3]

    In: The Eleventh International Conference on Learning Representations (2022) Elhage,N.,Hume,T.,Olsson,C.,Schiefer,N.,Henighan,T.,Kravec,S.,Hatfield-

    Cundy, C., Hutter, M., Legg, S., Veness, J., et al.: Neural networks and the chomsky hierarchy. In: The Eleventh International Conference on Learning Representations (2022) Elhage,N.,Hume,T.,Olsson,C.,Schiefer,N.,Henighan,T.,Kravec,S.,Hatfield-

  4. [4]

    arXiv preprint arXiv:2209.10652 (2022)

    Dodds, Z., Lasenby, R., Drain, D., Chen, C., et al.: Toy models of superposi- tion. arXiv preprint arXiv:2209.10652 (2022)

  5. [5]

    Reorganizing attention-space geometry with expressive attention

    Gros, C.: Reorganizing attention-space geometry with expressive attention. arXiv preprint arXiv:2407.18601 (2024)

  6. [6]

    arXiv preprint arXiv:2312.00752 (2023)

    Gu, A., Dao, T.: Mamba: Linear-time sequence modeling with selective state spaces. arXiv preprint arXiv:2312.00752 (2023)

  7. [7]

    In: 2021 IEEE International Conference on Robotics and Automation (ICRA), pp

    Levine, S.: Reset-free reinforcement learning via multi-task learning: Learning dexterous manipulation behaviors without human intervention. In: 2021 IEEE International Conference on Robotics and Automation (ICRA), pp. 6664– 6671, IEEE (2021)

  8. [8]

    arXiv preprint arXiv:2203.15556 (2022)

    Hoffmann, J., Borgeaud, S., Mensch, A., Buchatskaya, E., Cai, T., Ruther- ford, E., Casas, D.d.L., Hendricks, L.A., Welbl, J., Clark, A., et al.: Train- ing compute-optimal large language models. arXiv preprint arXiv:2203.15556 (2022)

Show all 25 references
  1. [9]

    In: Findings of the Association for Computational Linguistics: EMNLP 2022, pp

    Ivgi, M., Carmon, Y., Berant, J.: Scaling laws under the microscope: Predict- ing transformer performance from small scale experiments. In: Findings of the Association for Computational Linguistics: EMNLP 2022, pp. 7354–7371 (2022)

  2. [10]

    arXiv preprint arXiv:2001.08361 (2020) Katharopoulos,A.,Vyas,A.,Pappas,N.,Fleuret,F.:Transformersarernns:Fast autoregressive transformers with linear attention

    Gray, S., Radford, A., Wu, J., Amodei, D.: Scaling laws for neural language models. arXiv preprint arXiv:2001.08361 (2020) Katharopoulos,A.,Vyas,A.,Pappas,N.,Fleuret,F.:Transformersarernns:Fast autoregressive transformers with linear attention. In: International conference on ...

  3. [11]

    Advances in Neural Information Processing Systems36 (2024) Kumar,V.,Shah,R.,Zhou,G.,Moens,V.,Caggiano,V.,Gupta,A.,Rajeswaran, A.: Robohive: A unified framework for robot learning

    Knight, P., Duan, R.: Multi-task learning with summary statistics. Advances in Neural Information Processing Systems36 (2024) Kumar,V.,Shah,R.,Zhou,G.,Moens,V.,Caggiano,V.,Gupta,A.,Rajeswaran, A.: Robohive: A unified framework for robot learning. Advances in Neural Information...

  4. [12]

    Transactions of the Association for Computational Linguistics12, 157–173 (2024)

    Liu, N.F., Lin, K., Hewitt, J., Paranjape, A., Bevilacqua, M., Petroni, F., Liang, P.: Lost in the middle: How language models use long contexts. Transactions of the Association for Computational Linguistics12, 157–173 (2024)

  5. [13]

    arXiv preprint arXiv:2307.06435 (2023)

    Barnes, N., Mian, A.: A comprehensive overview of large language models. arXiv preprint arXiv:2307.06435 (2023)

  6. [14]

    In: Deep Reinforcement Learning Workshop NeurIPS 2022 (2022)

    Neumann, O., Gros, C.: Scaling laws for a multi-agent reinforcement learning model. In: Deep Reinforcement Learning Workshop NeurIPS 2022 (2022)

  7. [15]

    arXiv preprint arXiv:2412.11979 (2024)

    Neumann, O., Gros, C.: Alphazero neural scaling and zipf’s law: a tale of board games and power laws. arXiv preprint arXiv:2412.11979 (2024)

  8. [16]

    Press, O., Smith, N.A., Lewis, M.: Train short, test long: Attention with linear biases enables input length extrapolation. arXiv preprint arXiv:2108.12409 (2021) Sándor, B., Nowak, M., Koglin, T., Martin, L., Gros, C.: Kick control: using the attracting states arising within ...

  9. [17]

    The International Journal of Robotics Research42(13), 1133–1184 (2023)

    Saveriano, M., Abu-Dakka, F.J., Kramberger, A., Peternel, L.: Dynamic move- ment primitives in robotics: A tutorial survey. The International Journal of Robotics Research42(13), 1133–1184 (2023)

  10. [18]

    arXiv preprint arXiv:2406.16690 (2024)

    Shen, X., Li, D., Leng, R., Qin, Z., Sun, W., Zhong, Y.: Scaling laws for linear complexity language models. arXiv preprint arXiv:2406.16690 (2024)

  11. [19]

    Transactions of the Association for Computational Linguistics12, 543–561 (2024)

    Strobl, L., Merrill, W., Weiss, G., Chiang, D., Angluin, D.: What formal lan- guages can transformers express? a survey. Transactions of the Association for Computational Linguistics12, 543–561 (2024)

  12. [20]

    arXiv preprint arXiv:2011.04006 (2020)

    Tay, Y., Dehghani, M., Abnar, S., Shen, Y., Bahri, D., Pham, P., Rao, J., Yang, L., Ruder, S., Metzler, D.: Long range arena: A benchmark for efficient trans- formers. arXiv preprint arXiv:2011.04006 (2020)

  13. [21]

    arXiv preprint arXiv:2405.15618 (2024)

    Tong, W.L., Pehlevan, C.: Mlps learn in-context. arXiv preprint arXiv:2405.15618 (2024)

  14. [22]

    Advances in neural infor- mation processing systems30 (2017)

    Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A.N., Kaiser, Ł., Polosukhin, I.: Attention is all you need. Advances in neural infor- mation processing systems30 (2017)

  15. [23]

    arXiv preprint arXiv:2006.04768 (2020) Wu,H.,Wu,J.,Xu,J.,Wang,J.,Long,M.:Flowformer:Linearizingtransformers with conservation flows

    Wang, S., Li, B.Z., Khabsa, M., Fang, H., Ma, H.: Linformer: Self-attention with linear complexity. arXiv preprint arXiv:2006.04768 (2020) Wu,H.,Wu,J.,Xu,J.,Wang,J.,Long,M.:Flowformer:Linearizingtransformers with conservation flows. arXiv preprint arXiv:2202.06258 (2022)

  16. [24]

    arXiv preprint arXiv:2206.04301 (2022)

    Zhang, Y., Backurs, A., Bubeck, S., Eldan, R., Gunasekar, S., Wagner, T.: Unveiling transformers with lego: a synthetic reasoning task. arXiv preprint arXiv:2206.04301 (2022)

  17. [25]

    IEEE transactions on knowledge and data engineering34(12), 5586–5609 (2021)

    Zhang, Y., Yang, Q.: A survey on multi-task learning. IEEE transactions on knowledge and data engineering34(12), 5586–5609 (2021)

Pith tools

Reviewed August 5, 2026 · model on record in the stance chip above.