Gil Tabak
Authored Publications
Sort By
Preview abstract
Linear attention compresses context into a fixed-size key-value state for linear-time training and constant-memory inference. Gated delta-rule models decay this state and erase its prediction for an incoming key before writing a new value; Gated DeltaNet-2 (GDN-2) generalizes this with channel-wise decay and decoupled, full-rank erase and write gates. However, we show that state updates are far less sensitive to erasure than to writes: key activation and unit normalization concentrate key energy onto a sparse subset of channels, whereas value writes scale all channels, so the erase gate can be parameterized at lower rank. Furthermore, slow-decaying state rows accumulate much larger norms than fast rows, dominating unweighted readout and forcing stronger write-time erasure to suppress interference. Guided by this analysis, we introduce Muir, which fuses all gates into a sensitivity-allocated low-rank projection, re-weights state rows with a matched-filter readout gate, and reinvests the saved parameters in a wider feed-forward network. Because standard chunkwise algorithms become numerically unstable at the large chunk sizes needed to saturate wide matrix units, we design a parallel update for stable, efficient training. Across 166M–1.34B parameters, Muir lowers perplexity by 1.8–2.0% over GDN-2 and 2.5–3.6% over Transformer, raises commonsense accuracy by up to 1.9 points, and improves extractive recall over GDN-2. At 1.34B, Muir outperforms both baselines on multi-hop and multi-needle recall, extrapolates 2,560x from its 4K training length to 10M-token streams without perplexity drift, and maintains near-100% passkey retrieval to 1M tokens with static readout. Muir achieves these gains while matching GDN-2 throughput, reaching its 100B-token validation loss 1.2x faster, and training 2.7x faster than Transformer at 32K.
View details
EQuARX: Efficient Quantized AllReduce in XLA for Distributed Machine Learning Acceleration
Ibrahim Ahmed
Clemens Schaefer
Denis Vnukov
Zenong Zhang
Toli Yevtushenko
2025
Preview abstract
Many digital advertisers continue to rely on attribution models to estimate the effectiveness of their marketing spend, allocate budget, and guide bidding decisions for real time auctions. The work described in this paper builds on previous efforts to better understand the capabilities and limitations of attribution models using simulated path data with experiment-based ground truth. While previous efforts were based on a generic specification of user path characteristics (e.g., ad channels considered, observed events included, and the transition rates between observed events), here we generalize the process to include a pre-analysis optimization step that matches the characteristics of the simulated path data with a set of reference path data from a particular advertiser. An attribution model analysis conducted with path-matched data is more relevant and applicable to an advertiser than generic path data. We demonstrate this path-fitting process using data from Booking.com. The simulated matched paths are used to demonstrate a few key capabilities and limitations for several position-based attribution models.
View details
Preview abstract
Profiling cellular phenotypes from microscopic imaging can provide meaningful biological information resulting from various factors affecting the cells. One motivating application is drug development: morphological cell features can be captured from images, from which similarities between different drug compounds applied at different doses can be quantified. The general approach is to find a function mapping the images to an embedding space of manageable dimensionality whose geometry captures relevant features of the input images. An important known issue for such methods is separating relevant biological signal from nuisance variation. For example, the embedding vectors tend to be more correlated for cells that were cultured and imaged during the same week than for those from different weeks, despite having identical drug compounds applied in both cases. In this case, the particular batch in which a set of experiments were conducted constitutes the domain of the data; an ideal set of image embeddings should contain only the relevant biological information (e.g., drug effects). We develop a general framework for adjusting the image embeddings in order to “forget” domain-specific information while preserving relevant biological information. To achieve this, we minimize a loss function based on distances between marginal distributions (such as the Wasserstein distance) of embeddings across domains for each replicated treatment. For the dataset we present results with, the only replicated treatment happens to be the negative control treatment, for which we do not expect any treatment-induced cell morphology changes. We find that for our transformed embeddings (i) the underlying geometric structure is not only preserved but the embeddings also carry improved biological signal; and (ii) less domain-specific information is present.
View details
