Official Review · ICLR 2025

The paper investigates the application of Memory-Efficient transformers under the Algorithm Distillation for that leverages in-context learning abilities of autoregressive transformers for better generalization. Specifically, the authors propose a benchmark suite that efficiently assesses ICL capabilities within AD. The benchmark is demonstrated by testing on various memory-efficient transformers. The ERNIE-Docs (ED), a variant of Transformer-XL structure, is highlighted for its competitive performance with reduced memory requirements. The benchmark suite covers diverse reinforcement learning tasks, including scenarios with sparse and dense rewards, requiring credit assignment and exploration-exploitation balancing. The suite utilizes both JAX for efficient parallel execution and compatibility with PyTorch. Experimental results reveal that ED performs comparably to standard AD but with a lower memory footprint, showing promise for resource-constrained environments. Notably, the study identifies key factors like positional encoding, model size, attention settings, and memory capacity as influential to the ICL performance of ED. + Novel Application of Efficient Transformers in In-Context Reinforcement Learning: The paper introduces a unique connection by leveraging Efficient Transformer (ET) architectures, specifically ERNIE-Docs, within Algorithm Distillation (AD) for in-context reinforcement learning (ICL). + Comprehensive Benchmarking Suite for Memory-Efficient Algorithm Distillation: The paper’s proposed benchmark suite is tailored to evaluate ICL in AD using memory-efficient transformers. The benchmark suite facilitates thorough testing of memory-efficient AD methods across various reinforcement learning challenges, fulfilling the proposed property for meta RL benchmarking. + Systematic Ablation Studies and Insights into Hyperparameters: The paper provides detailed ablation studies, shedding light on how positional encoding, model size, attention settings, and recurrence memory capacity impact ICL performance. + Lack of Benchmarking Comparisons: Although the paper’s main contribution is a new benchmarking suite, it does not include comparisons with existing benchmarks. This omission makes it challenging to assess whether the proposed benchmark genuinely outperforms or offers advantages over established alternatives in terms of efficiency, coverage, or flexibility. + Limited Range of Tested Methods: As a benchmarking study, the paper evaluates a relatively narrow set of methods, primarily focusing on ERNIE-Docs and Transformer-XL. This limited scope overlooks a wide array of memory-efficient transformers in the literature, reducing the generalizability of the benchmark results and their applicability across diverse models. + Engineering Efforts Framed as Research Contributions: While the paper uses JAX for parallelization to achieve computational efficiency, this is more of an engineering choice than a novel research contribution. Clea