← Search

Alberto Bietti

35 accepted papers

2026

Learning to Recall with Transformers Beyond Orthogonal Embeddings

ICLR 2026poster

Modern large language models (LLMs) excel at tasks that require storing and retrieving knowledge, such as factual recall and question answering. Transformers are central to this capability, thanks to their ability to encode information during training and retrieve it at inference. Existing theoretic…

Cited by 0SourceScholar
2026

Protein Design with Agent Rosetta: A Case Study for Specialized Scientific Agents

ICML 2026poster

Large language models (LLMs) are capable of emulating reasoning and using tools, creating opportunities for autonomous agents that execute complex scientific tasks. Protein design provides a natural testbed: although machine learning (ML) methods achieve strong results, these are largely restricted …

Cited by 0SourceScholar
2026

Walrus: A Cross-domain Foundation Model for Continuum Dynamics

ICML 2026spotlight

Foundation models have transformed machine learning for language and vision, but achieving comparable impact in physical simulation remains a challenge. Data heterogeneity and unstable long-term dynamics inhibit learning from sufficiently diverse dynamics, while varying resolutions and dimensionalit…

Cited by 0SourceScholar
2025

AION-1: Omnimodal Foundation Model for Astronomical Sciences

NeurIPS 2025poster

While foundation models have shown promise across a variety of fields, astronomy lacks a unified framework for joint modeling across its highly diverse data modalities. In this paper, we present AION-1, the first large-scale multimodal foundation family of models for astronomy. AION-1 enables arbitr…

Cited by 0SourceScholar
2025

Distributional Associations vs In-Context Reasoning: A Study of Feed-forward and Attention Layers

ICLR 2025poster

Large language models have been successful at tasks involving basic forms of in-context reasoning, such as generating coherent language, as well as storing vast amounts of knowledge. At the core of the Transformer architecture behind such models are feed-forward and attention layers, which are often…

Cited by 0SourcePDFScholar
2025

Emergence of Linear Truth Encodings in Language Models

NeurIPS 2025poster

Recent probing studies reveal that large language models exhibit linear subspaces that separate true from false statements, yet the mechanism behind their emergence is unclear. We introduce a transparent, one-layer transformer toy model that reproduces such truth subspaces end-to-end and exposes one…

Cited by 0SourceScholar
2025

From Shortcut to Induction Head: How Data Diversity Shapes Algorithm Selection in Transformers

NeurIPS 2025spotlight

Transformers can implement both generalizable algorithms (e.g., induction heads) and simple positional shortcuts (e.g., memorizing fixed output positions). In this work, we study how the choice of pretraining data distribution steers a shallow transformer toward one behavior or the other. Focusing o…

Cited by 0SourceScholar
2025

In-Context Denoising with One-Layer Transformers: Connections between Attention and Associative Memory Retrieval

ICML 2025oral

We introduce in-context denoising, a task that refines the connection between attention-based architectures and dense associative memory (DAM) networks, also known as modern Hopfield networks. Using a Bayesian framework, we show theoretically and empirically that certain restricted denoising problem…

Cited by 0SourcePDFScholar
2025

Predicting partially observable dynamical systems via diffusion models with a multiscale inference scheme

NeurIPS 2025poster

Conditional diffusion models provide a natural framework for probabilistic prediction of dynamical systems and have been successfully applied to fluid dynamics and weather prediction. However, in many settings, the available information at a given time represents only a small fraction of what is nee…

Cited by 0SourceScholar
2025

Understanding Factual Recall in Transformers via Associative Memories

ICLR 2025spotlight

Large language models have demonstrated an impressive ability to perform factual recall. Prior work has found that transformers trained on factual recall tasks can store information at a rate proportional to their parameter count. In our work, we show that shallow transformers can use a combination…

Cited by 7SourcePDFScholar
2024

Heavy-Tailed Class Imbalance and Why Adam Outperforms Gradient Descent on Language Models

NeurIPS 2024spotlight

Adam has been shown to outperform gradient descent on large language models by a larger margin than on other tasks, but it is unclear why. We show that a key factor in this performance gap is the heavy-tailed class imbalance found in language tasks. When trained with gradient descent, the loss of in…

Cited by 32SourcePDFScholar
2024

Multiple Physics Pretraining for Spatiotemporal Surrogate Models

NeurIPS 2024poster

We introduce multiple physics pretraining (MPP), an autoregressive task-agnostic pretraining approach for physical surrogate modeling of spatiotemporal systems with transformers. In MPP, rather than training one model on a specific physical system, we train a backbone model to predict the dynamics o…

Cited by 3SourcePDFScholar
2023

Birth of a Transformer: A Memory Viewpoint

NeurIPS 2023spotlight

Large language models based on transformers have achieved great empirical successes. However, as they are deployed more widely, there is a growing need to better understand their internal mechanisms in order to make them more reliable. These models appear to store vast amounts of knowledge from thei…

Cited by 86SourcePDFScholar
2023

The SSL Interplay: Augmentations, Inductive Bias, and Generalization

ICML 2023poster

Self-supervised learning (SSL) has emerged as a powerful framework to learn representations from raw data without supervision. Yet in practice, engineers face issues such as instability in tuning optimizers and collapse of representations during training. Such challenges motivate the need for a theo…

Cited by 41SourcePDFScholar
2022

Efficient Kernelized UCB for Contextual Bandits

AISTATS 2022poster

In this paper, we tackle the computational efficiency of kernelized UCB algorithms in contextual bandits. While standard methods require a $\mathcal{O}(CT^3)$ complexity where $T$ is the horizon and the constant $C$ is related to optimizing the UCB rule, we propose an efficient contextual algorithm…

Cited by 24SourcePDFScholar
2022

Learning single-index models with shallow neural networks

NeurIPS 2022accept

Single-index models are a class of functions given by an unknown univariate ``link'' function applied to an unknown one-dimensional projection of the input. These models are particularly relevant in high dimension, when the data might present low-dimensional structure that learning algorithms should…

Cited by 102SourcePDFScholar
2022

Personalization Improves Privacy-Accuracy Tradeoffs in Federated Learning

ICML 2022spotlight

Large-scale machine learning systems often involve data distributed across a collection of users. Federated learning algorithms leverage this structure by communicating model updates to a central server, rather than entire datasets. In this paper, we study stochastic optimization algorithms for a pe…

2022

When does return-conditioned supervised learning work for offline reinforcement learning?

NeurIPS 2022accept

Several recent works have proposed a class of algorithms for the offline reinforcement learning (RL) problem that we will refer to as return-conditioned supervised learning (RCSL). RCSL algorithms learn the distribution of actions conditioned on both the state and the return of the trajectory. Then…

2021

On Energy-Based Models with Overparametrized Shallow Neural Networks

ICML 2021oral

Energy-based models (EBMs) are a simple yet powerful framework for generative modeling. They are based on a trainable energy function which defines an associated Gibbs measure, and they can be trained and sampled from via well-established statistical tools, such as MCMC. Neural networks may be used…

2021

On the Universality of Graph Neural Networks on Large Random Graphs

NeurIPS 2021poster

We study the approximation power of Graph Neural Networks (GNNs) on latent position random graphs. In the large graph limit, GNNs are known to converge to certain ``continuous'' models known as c-GNNs, which directly enables a study of their approximation power on random graph models. In the absence…

2020

Convergence and Stability of Graph Convolutional Networks on Large Random Graphs

NeurIPS 2020spotlight

We study properties of Graph Convolutional Networks (GCNs) by analyzing their behavior on standard models of random graphs, where nodes are represented by random latent variables and edges are drawn according to a similarity kernel. This allows us to overcome the difficulties of dealing with discre…

2019

A Kernel Perspective for Regularizing Deep Neural Networks

ICML 2019oral

We propose a new point of view for regularizing deep neural networks by using the norm of a reproducing kernel Hilbert space (RKHS). Even though this norm cannot be computed, it admits upper and lower approximations leading to various practical strategies. Specifically, this perspective (i) provides…

2017

Stochastic Optimization with Variance Reduction for Infinite Datasets with Finite Sum Structure

NeurIPS 2017spotlight

Stochastic optimization algorithms with variance reduction have proven successful for minimizing large finite sums of functions. Unfortunately, these techniques are unable to deal with stochastic perturbations of input data, induced for example by data augmentation. In such cases, the objective is n…

2015

An online EM algorithm in hidden (semi-)Markov models for audio segmentation and clustering

ICASSP 2015accepted

Audio segmentation is an essential problem in many audio signal processing tasks, which tries to segment an audio signal into homogeneous chunks. Rather than separately finding change points and computing similarities between segments, we focus on joint segmentation and clustering, using the framewo…

Cited by 0SourceScholar