Balancing Gradient Frequencies Facilitates Inductive Inference in Algorithmic Reasoning

Thank you for introducing this interesting paper. However, we believe that the suggested paper and our work are different in multiple aspects.

The paper you mentioned, “Frequency Shortcut Learning in Neural Networks”, focuses on analyzing input data (images) in the Fourier domain and identifying specific frequencies that act as shortcuts. Indeed, transforming input data into the frequency domain to analyze spurious correlations or components related to generalization is a well-established area of research, particularly in the vision domain (e.g., High-frequency Component Helps Explain the Generalization of Convolutional Neural Networks, CVPR ‘20).

In contrast, our work takes a fundamentally different approach by linking the frequency of gradients to the optimization performance in algorithmic reasoning (AR) tasks. Specifically, we introduce temporal frequency filtering during the learning process to guide the model toward learning generalizing hypotheses. This is a novel perspective that has not been explored in prior work.

Additionally, we highlight in Section 5.1 that optimizers like SAM, which have demonstrated strong generalization performance in vision tasks, perform poorly in AR tasks. This result further underscores the unique nature of AR tasks and the importance of studying optimization dynamics specifically tailored to these problems.

We hope this clarifies the distinction between the contributions of our work and the referenced paper, and we are grateful for the opportunity to elaborate on these points.

Thank you for your thoughtful and detailed feedback. We appreciate the opportunity to clarify and strengthen the discussion around the definition of “shortcuts” in our paper.

Regarding the definition of shortcuts

We acknowledge that “shortcuts” have been broadly defined in prior work, often in the context of classification tasks, as decision rules leveraging spurious correlations specific to a dataset, rather than semantically meaningful task-related cues. In AR, however, no exact definition of shortcuts currently exists, but previous works in AR (What Algorithms can Transformers Learn? A Study in Length Generalization, ICLR ’24 / Transformers Learn Shortcuts to Automata, ICLR ’22 Oral) found that “shortcuts” correspond to non-generalizing hypotheses in reasoning tasks.

As the reviewer notes, “shortcut learning” in classification tasks typically involve using easily learnable but dataset-specific features, rather than complex core features, due to the simplicity bias in Deep Neural Networks (DNNs). On the contrary, performing inductive inference to solve AR tasks requires models to generalize beyond the training distribution, distinguishing it from typical classification problems. The simplicity bias in AR usually occurs because DNNs learn non-generalizing hypotheses that are compatible with the training data but fail to extrapolate to out-of-distribution (OOD) tasks. Therefore, we opted to define the non-generalizing hypotheses as shortcuts specifically in the context of AR.

Below, we clarify distinctions between AR and conventional classification problems that necessitate a separate definition of shortcuts for AR:

1. Exhaustion of Training Data Information: As mentioned in our paper, AR tasks often result in near-perfect training accuracy with the Adam optimizer, highlighting that the model utilizes all available information in the training data. However, the inductive leap required for OOD generalization depends on the ability to infer general logic rather than settle for non-generalizing patterns.
2. The Role of Inductive Bias: In AR, the challenge stems not from choosing between spurious and meaningful features but from selecting the “correct” generalization among infinite plausible hypotheses. For example, predicting the continuation of the sequence “1, 3, 5, 7” could yield numerous valid hypotheses, but general logic identifies “9” as the most parsimonious continuation.
3. Non-Generalizing Hypotheses as Shortcuts: In this context, we term “shortcuts” as those hypotheses that fail to generalize OOD, despite being valid within the training distribution. These shortcuts may not necessarily be simpler than the “general logic” but are misleading in the sense that they align with the training data but not the target task.

On contents of lines 191–193

You pointed out that the statement “Shortcuts can be interpreted as learning non-generalizing correlations between inputs and outputs that are valid in the training distribution but do not hold in the target distribution” is not fully accurate as it depends on the target distribution. We agree with this nuance and appreciate the chance to refine our discussion. Our intention was to highlight that in AR tasks, the target distribution is always OOD (lines 30-32, 122-125, and 189-191) making non-generalizing correlations misleading for the intended task. However, as you rightly note, if the target distribution remains in-distribution, such correlations could still hold.

Revised Definition Proposal

We sincerely apologize that the reviewer found our choice of terminology misleading, despite our efforts to draw a clear distinction between AR and conventional classification. Taking into account both your feedback and the general understanding of shortcuts in the literature, we propose refining our definition as follows:

Definition: In algorithmic reasoning (AR), “shortcuts” refer to non-generalizing hypotheses that align with training data correlations but fail to extrapolate to out-of-distribution (OOD) tasks, which often define the target of AR problems. Unlike shortcuts in classification tasks—where they often involve simplicity bias favoring spurious correlations—shortcuts in AR tasks stem from the model’s inability to infer general logic among multiple plausible hypotheses fitting the training data.

We believe this revision better aligns our work with the existing literature while addressing the unique challenges posed by AR tasks. We hope this explanation clarifies our perspective and makes our contributions to the understanding of shortcut learning more compelling.

Thank you for your response. However, I still have concerns with this paper.

The authors’ new definition, ““shortcuts” refer to non-generalizing hypotheses that align with training data correlations but fail to extrapolate to out-of-distribution (OOD) tasks” is not correct. There are multiple factors that hurt (OOD) generalization, and shortcut learning is one of them. That is, non-generalizing hypotheses can include shortcuts, but shortcuts solely cannot represent non-generalizing hypotheses. Also, if the authors would argue, “…..shortcuts may not necessarily be simpler than the “general logic.””, which is not true in classification tasks, they should not use the term “shortcut” learning in this different context which alters and misleads the well-established meaning in the research community and literature. I guess, with my best, what authors try to say is one type of bias, but not shortcut learning – even the authors emphasize that it is different from the one defined in classification tasks. My concern is that this altering terminology or misuse of terminology will bring a lot of confusion to the community. Also, this creates a fundamental doubt about the authors' understanding of the phenomenon.

Thank you for your active engagement in the discussion period.

Please refer to Shortcut Learning in Deep Neural Networks, Nature Machine Intelligence ‘20 [1]. Figure 3 states that “Among the solutions that solve the training data, only some generalize to an i.i.d. test set. Among those solutions, shortcuts fail to generalize to different data (o.o.d. test sets), but the intended solution does generalize.” Therefore, we believe our definition of shortcuts is not incorrect.

Furthermore, the introduction of the above Nature paper states that “various phenomena [referring to learning under covariate shift, anti-causal learning, dataset bias, etc.] can be collectively termed shortcuts.” As stated in our initial response, the term “shortcuts” is flexibly used in AR [2,3,4], even if it does not strictly align with the the reviewer’s perspective. As such, we believe it is a stretch to argue that our use of the term “shortcuts” would alter or mislead its well-established meaning in the research community and literature.

[1] Geirhos, Robert, et al. "Shortcut learning in deep neural networks." Nature Machine Intelligence ’20

[2] Zhou, Yongchao, et al. "Transformers Can Achieve Length Generalization But Not Robustly." ICLR ‘24 - Workshop

[3] Liu, Bingbin, et al. "Transformers Learn Shortcuts to Automata." ICLR ‘22  

[4] Murty, Shikhar, et al. "Pushdown Layers: Encoding Recursive Structure in Transformer Language Models." EMNLP ‘23

Thank you for your considerate comments and for recognizing the novel insights and broad applicability of BGF in improving generalization for AR tasks. As you mentioned, our work is the first to approach emerging AR tasks from the perspective of optimization. To reflect your comments, we added new comparison results between BGF and recent optimizers from ‘23 and ‘24 (More comparison with recent optimizers). These additional results further highlight the uniqueness of AR tasks and the superior performance of BGF in addressing their challenges. Furthermore, at the request of Reviewer xigM, we verified BGF on decoder-only transformers and a real-world “composition” task (Additional analyses and results). The success of BGF at improving the performance of decoder-only transformers on this real-world reasoning task indicates that BGF can indeed be used to improve the reasoning abilities of Deep Neural Networks.

More comparison with recent optimizers

In Table 2 of the main paper, we compared BGF to SAM, SWAD, and LPF-SGD, and In Appendix 5, we compared BGF to AdamW. We additionally compare BGF to 2 recent optimization techniques: Marg-Ctrl loss and FSAM optimizer.

• Marg-Ctrl Loss (NeurIPS '23) [1] : The authors of the Marg-Ctrl Loss paper observed that the default Empirical Risk Minimization loss prefers solutions that maximize the margin. Inspired by this observation, the Marg-Ctrl loss was developed to enforce uniform-margin solutions. It showed notable ability to mitigate shortcuts in vision and language tasks. Below, we compare BGF to Marg-Ctrl loss. Because Marg-Ctrl loss was defined for a single-class classification problem, the comparative study is conducted on seven AR tasks with output length = 1. Like the rest of optimizers in the main paper, Marg-Ctrl falls short of BGF in all tasks.

• FSAM Optimizer (CVPR '24) [2]: The FSAM paper first decomposed the adversarial perturbations induced by SAM to analyze which gradient components play a crucial role in improving generalization. Their analysis revealed the importance of stochastic, rather than full, gradient noise components in promoting generalization, and thus, FSAM advances the original SAM optimizer by removing the full gradient noise components, which are detrimental to the mini-batch-wise optimization process. It demonstrates state-of-the-art generalization performance on the vision task. Additional comparative experiments on more seeds, architectures, and tasks are running and will be updated throughout the discussion period.
  • The table below compares the average of best accuracies since not all seeds of the FSAM experiments have finished running yet.

[1] Puli, Aahlad Manas, et al. "Don’t blame dataset shift! shortcut learning due to gradients and cross entropy."

[2] Li, Tao, et al. "Friendly sharpness-aware minimization."

Hyperparameter analysis on $\lambda$ and gradient balancing factors ($\alpha$ and $\beta$)

• Effect of changing $\lambda$: We used a fixed window size of 100 for all of our experiments. This value was chosen based on the gradient analysis shown in Figure 3, which demonstrated that the gradient patterns of both the training and generalization phases begin to change within the frequency range of 10^(-2)Hz and 10^(-1.5)Hz. A window size of 100 was appropriately chosen to suppress gradients in higher frequencies than this range.

  In the table below, we study the effect of changing the smoothing factor of BGF-ema, which is equivalent to altering $\lambda$ of the original BGF. The larger the smoothing factor is, the larger $\lambda$ becomes. The results are reported in terms of the best accuracy averaged over tasks and 3 seeds. This ablation study was conducted BGF-ema instead of the original BGF due to the time constraint of the discussion period. According to the results, 0.98 appears to be the apposite smoothing factor. Also, BGF outperforms Adam at smoothing factors ranging from 0.95 to 0.99.

• Effect of changing gradient balancing factors: For [$\alpha$, $\beta$] combination, the following combinations were tested: [0.95, 0.05], [0.90, 0.10], [0.80, 0.20], [0.70, 0.30]. Hyperparameter sweep was conducted using sequences of lengths 50~59, equivalent to mild OOD sequences. We used 128 samples per sequence length. The average of results from a total of 1280 samples (=128 x 10) was used to determine the best hyperparameter for each optimizer. To evidence that the results on mild OOD sequences are indeed representative of performance on longer, more extreme OOD sequences, hyperparameter sweep results and the final test accuracy on sequence length=100 are juxtaposed below.

1. RNN
   • Average validation accuracy on sequences of lengths 50 to 59 |Optim.|Acc.| |-|-| |Adam|0.852| |BGF[0.95, 0.05]|0.800| |BGF[0.90, 0.10]|0.829| |BGF[0.80, 0.20]|0.903| |BGF[0.70, 0.30]|0.917|
   • Test accuracy on sequences of length=100 |Optim.|Acc.| |-|-| |Adam|0.738| |BGF[0.95, 0.05]|0.734| |BGF[0.90, 0.10]|0.804| |BGF[0.80, 0.20]|0.831| |BGF[0.70, 0.30]|0.867|
2. LSTM (Note that BGF yields 1.000 accuracy on two hyperparameter settings. In such cases, we selected the one that converged fatser to 0.9 according to the number of training steps)
   • Average validation accuracy on sequences of lengths 50 to 59 |Optim.|Acc.| |-|-| |Adam|0.998| |BGF[0.95, 0.05]|1.000| |BGF[0.90, 0.10]|0.999| |BGF[0.80, 0.20]|1.000| |BGF[0.70, 0.30]|0.999|
   • Test accuracy on sequences of length=100 |Optim.|Acc.| |-|-| |Adam|0.957| |BGF[0.95, 0.05]|0.968| |BGF[0.90, 0.10]|0.965| |BGF[0.80, 0.20]|0.989| |BGF[0.70, 0.30]|0.986|
   • Number of training steps to 0.9 |Optim.|Acc.| |-|-| |Adam|1248| |BGF[0.95, 0.05]|911| |BGF[0.90, 0.10]|663| |BGF[0.80, 0.20]|707| |BGF[0.70, 0.30]|910|

Gradient analysis on different layers of DNN architectures

• We used a 4-layer LSTM model to analyze the gradients from different DNN layers. This experiment was conducted on the solve equation task. Please refer to Figure R3 of the revised paper for the visualization results. In early layers (Figure R3 (a)), the training and generalization phase gradients show a clear split between the two, while in later layers (Figure R3 (b)), this gap between the two gradients is visibly reduced. From this result, we conjecture that early layers could benefit more from suppressing high-frequency gradients.

• We additionally analyzed the gradients from different layers of the Transformer model on the composition task - one of the most critical reasoning tasks in the real world (as discussed in Grokked Transformers are Implicit Reasoners: A Mechanistic Journey to the Edge of Generalization, NeurIPS '24). Please refer to Figure R2 of the revised paper for the visualization results. We observe a similar tendency to the above in this experiment; the discrepancy between training and generalization phase gradients is visibly larger in early layer gradients. This result again implies that early layers could benefit more from suppressing high-frequency gradients.

Thank you for your insightful comments and for recognizing the novelty and effectiveness of our approach in improving OOD generalization for AR tasks through a new optimizer. As you suggested, we added new comparison results between BGF and recent optimizers from ‘23 and ‘24 (More comparison with recent optimizers). These additional results further highlight the uniqueness of AR tasks and the superior performance of BGF in addressing their challenges. Furthermore, at the request of Reviewer xigM, we verified BGF on decoder-only transformers and a real-world “composition” task (Additional analyses and results). The success of BGF at improving the performance of decoder-only transformers on this real-world reasoning task indicates that BGF can indeed be used to improve the reasoning abilities of Deep Neural Networks.

Theoretical analysis

We humbly agree with the reviewer in that our paper can be improved with a formal theoretical analysis of BGF. Unfortunately, Algorithmic Reasoning is an emerging and underexplored field in and of itself, and research on theoretical frameworks for analyzing the properties of AR remains to be done. We sincerely wished to accompany the proposed optimizer with more theoretical analyses, such as its convergence guarantees, but the lack of theoretical tools for studying AR left us with no other option but to strengthen our claim with extensive empirical evidence and analyses. We believe developing a theoretical framework to better understanding reasoning capabilities Deep Neural Networks would be an impactful future direction of research and will continue our endeavors to comprehend DNNs’ reasoning behaviors.

More comparison with recent optimizers

In Table 2 of the main paper, we compared BGF to SAM, SWAD, and LPF-SGD, and In Appendix 5, we compared BGF to AdamW. We additionally compare BGF to 2 recent optimization techniques: Marg-Ctrl loss and FSAM optimizer.

• Marg-Ctrl Loss (NeurIPS '23) [1] : The authors of the Marg-Ctrl Loss paper observed that the default Empirical Risk Minimization loss prefers solutions that maximize the margin. Inspired by this observation, the Marg-Ctrl loss was developed to enforce uniform-margin solutions. It showed notable ability to mitigate shortcuts in vision and language tasks. Below, we compare BGF to Marg-Ctrl loss. Because Marg-Ctrl loss was defined for a single-class classification problem, the comparative study is conducted on seven AR tasks with output length = 1. Like the rest of optimizers in the main paper, Marg-Ctrl falls short of BGF in all tasks.

• FSAM Optimizer (CVPR '24) [2]: The FSAM paper first decomposed the adversarial perturbations induced by SAM to analyze which gradient components play a crucial role in improving generalization. Their analysis revealed the importance of stochastic, rather than full, gradient noise components in promoting generalization, and thus, FSAM advances the original SAM optimizer by removing the full gradient noise components, which are detrimental to the mini-batch-wise optimization process. It demonstrates state-of-the-art generalization performance on the vision task. Additional comparative experiments on more seeds, architectures, and tasks are running and will be updated throughout the discussion period.
  • The table below compares the average of best accuracies since not all seeds of the FSAM experiments have finished running yet.

[1] Puli, Aahlad Manas, et al. "Don’t blame dataset shift! shortcut learning due to gradients and cross entropy."

[2] Li, Tao, et al. "Friendly sharpness-aware minimization."

Thank you for your thoughtful comments and for recognizing the strength of our paper in proposing a novel optimization method for improving generalization in AR tasks. As you mentioned, our work is the first to approach emerging AR tasks from the perspective of optimization. Following your suggestion, we 1) clarified our task descriptions and interpretation of experimental results and 2) studied the effect of changing $\lambda$. During the rebuttal process, we also added new experimental results to demonstrate the unique and practical advantages of BGF. First, at the request of Reviewers oa4Z and cFMM, we added new comparison results between BGF and recent optimizers from ‘23 and ‘24, which further highlight the uniqueness of AR tasks and the efficacy of BGF in addressing their challenges (Reviewers oa4Z, cFMM - More comparison with recent optimizers). Second, at the request of Reviewer xigM, we verified BGF on decoder-only transformers for a real-world “composition” task (Additional analyses and results). The transferability of BGF to decoder-only transformers on this real-world task indicates that BGF can indeed be used to improve the reasoning abilities of Deep Neural Networks.

Task descriptions

We apologize that task descriptions were not provided clearly in the main paper. We were trying to pack a lot of material within the 10-page constraint and had to move the detailed task descriptions to Appendix. In Pg. 4 of the revised paper, we included an example of how we plan to include task descriptions in the main paper (marked in red). Also, in Table A2 of the revised Appendix, we added a table with example input output pairs for each task to improve the readability of our work as a standalone paper.

Effect of changing $\lambda$

We used a fixed window size of 100 for all of our experiments. This value was chosen based on the gradient analysis shown in Figure 3, which demonstrated that the gradient patterns of both the training and generalization phases begin to change within the frequency range of 10^(-2)Hz and 10^(-1.5)Hz. A window size of 100 was appropriately chosen to suppress gradients in higher frequencies than this range.

In the table below, we study the effect of changing the smoothing factor of BGF-ema, which is equivalent to altering $\lambda$ of the original BGF. The larger the smoothing factor is, the larger $\lambda$ becomes. The results are reported in terms of the best accuracy averaged over tasks and 3 seeds. This ablation study was conducted BGF-ema instead of the original BGF due to the time constraint of the discussion period. According to the results, 0.98 appears to be the apposite smoothing factor. Also, BGF outperforms Adam at smoothing factors ranging from 0.95 to 0.99.

Why is the result of EMA lower than queue in Table 3 [Right]?

We note that on the LSTM model, BGF with queue and EMA both fail to achieve >90% accuracy, indicating that neither one of the implementations enabled LSTM to learn the generalizing logic of this task. Therefore, although BGF with queue attains slightly higher accuracy than that with EMA, it is difficult to ascertain that the queue-based implementation is superior to EMA-based implementation.

Additional analyses and results (Pg 26-27 of the revised paper)

In the revised manuscript, we 1) additionally verified BGF on a decoder-only transformer using the composition task—one of the most critical reasoning tasks in the real world (as discussed in Grokked Transformers are Implicit Reasoners: A Mechanistic Journey to the Edge of Generalization, NeurIPS '24) (Section R1 of the revised Appendix), and 2) conducted a layer-wise frequency analysis of gradient signals (Section R2 of the revised Appendix). The first experiment, performed with the decoder-only transformer on the composition task, corroborates that BGF can be utilized to promote learning of generalizing rules in more complex architectures and real-world reasoning tasks. The gradient analysis on the first experiment again reveals the discrepancy in training and generalization phase gradient signals. The second analysis indicates that early layer gradients are more likely to benefit from BGF, providing a further insight into the mechanism of BGF.