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 →
The pith
A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.
The reading
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.
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
- 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.
Editorial analysis
A structured set of objections, weighed in public.
Referee Report
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)
- [§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
- [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)
- [Abstract and §6] The abstract refers to 'extensive attention' and §6 refers to the 'cistransformer'; both should be 'expressive attention' and 'cisformer'.
- [§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.
- [§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.
- [§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.
- [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
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
free parameters (5)
- Learning rate and momentum =
0.02, 0.8
- Number of layers per model =
Transformer 60, cisformer 12, MLP 16, LSTM 2
- Context length Ncon =
24
- Batch size and epochs =
200, 8000
- Control token interval distribution =
6 +/- 3, equal probabilities
assumptions (5)
- domain assumption IARC task definition (Eq. 1-4): I/A/R/C operations as specified
- domain assumption Control tokens are stochastically interspersed with 6 +/- 3 spacing and equal probabilities
- domain assumption One-hot embedding with control tape (Eq. 5) and d = N + Ncontrol
- domain assumption ALiBi positional encoding (Press et al., 2021) is used for all attention models
- domain assumption Training with cross-entropy loss and SGD with momentum
invented entities (2)
-
Cisformer
-
Expressive attention (EA)
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
Reference graph
Works this paper leans on
-
[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)
work page 2024
-
[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)
work page 2024
-
[3]
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-
work page 2022
-
[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)
arXiv 2022
-
[5]
Reorganizing attention-space geometry with expressive attention
Gros, C.: Reorganizing attention-space geometry with expressive attention. arXiv preprint arXiv:2407.18601 (2024)
work page Pith review arXiv 2024
-
[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)
arXiv 2023
-
[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)
work page 2021
-
[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)
arXiv 2022
Show all 25 references
-
[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)
2022
-
[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 ...
2001 arXiv
-
[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...
2024
-
[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)
2024
-
[13]
arXiv preprint arXiv:2307.06435 (2023)
Barnes, N., Mian, A.: A comprehensive overview of large language models. arXiv preprint arXiv:2307.06435 (2023)
2023 arXiv
-
[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)
2022
-
[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)
2024
-
[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 ...
2021 arXiv
-
[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)
2023
-
[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)
2024 arXiv
-
[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)
2024
-
[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)
2011 arXiv
-
[21]
arXiv preprint arXiv:2405.15618 (2024)
Tong, W.L., Pehlevan, C.: Mlps learn in-context. arXiv preprint arXiv:2405.15618 (2024)
2024 arXiv
-
[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)
2017
-
[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)
2006 arXiv
-
[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)
2022 arXiv
-
[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)
2021
Reviewed August 5, 2026 · model on record in the stance chip above.
Discussion (0). Sign in to comment.