google-deepmind/regress-lm

google-deepmind/regress-lm

Description: Library for text-to-text regression, applicable to any input string representation and allows pretraining and fine-tuning over multiple regression tasks.

Language: Python

License: Apache-2.0

Stars: 344

Forks: 59

Open issues: 11

Created: 2025-06-09T15:08:53Z

Pushed: 2026-06-09T00:07:45Z

Default branch: main

Fork: no

Archived: no

README:

RegressLM: Easy Sequence-to-Sequence Regression


**Google Research Blog** | [Setup](#setup) | [Usage](#usage) | [Extended Usage](#extended_usage) | [Citing](#citing)

Core Contributors: Xingyou Song, Yash Akhauri, Jiyoun Ha, Bryan Lewandowski

Overview

RegressLM is a library for sequence-to-sequence regression, applicable to input string + image representations and allows pretraining and fine-tuning over multiple regression tasks.

Example Application: Directly regressing performance metrics from unstructured, textually represented system states from Google's massive compute clusters.

Setup

Get started by installing the core libraries:

pip install -e .

To run e.g. T5Gemma variants, install additional libraries:

pip install ".[extras]"

Usage

There are two main stages: inference and pretraining (optional but recommended).

Inference

The intended use-case is to import a RegressLM class, which can decode floating-point predictions from a given input, and also fine-tune against new data.

from regress_lm import core
from regress_lm import rlm

# Create RegressLM from scratch. Optionally, use `from_t5gemma_encoder`.
reg_lm = rlm.RegressLM.from_scratch(max_input_len=2048)

# Example (x,y) pairs, which can be fine-tuned against.
examples = [core.Example(x='hello', y=0.3), core.Example(x='world', y=-0.3)]
reg_lm.fine_tune(examples)

# Query inputs.
query1, query2 = core.ExampleInput(x='hi'), core.ExampleInput(x='bye')
samples1, samples2 = reg_lm.sample([query1, query2], num_samples=128)

Pretraining

To produce better initial checkpoints for transfer learning, we recommend the user pretrains over large amounts of their own training data. Example pseudocode with PyTorch:

from regress_lm.pytorch import model as model_lib
from regress_lm.pytorch import training

model = model_lib.PyTorchModelConfig(...).make_model()
trainer = training.Trainer(model, optimizer_factory, train_dataset, ...)
for batch in trainer.train_dl:
train_metrics = trainer.run_train_step(batch)

Boosting Performance and Extended Applications

Below, we describe ways to improve performance and extended applications, using lower level API.

Train Custom Vocabulary

You can generate a custom vocabulary, trained on an offline corpus of data mydata.txt:

encoder_vocab = SentencePieceVocab.from_corpus(corpus_path='mydata.txt', vocab_size=1024)

Larger Sizes

Larger model sizes may increase performance, although with more computational cost:

config = PyTorchModelConfig(architecture_kwargs=dict(num_encoder_layers=12, num_decoder_layers=12))

Multi-objective Support

The RLM can decode a concatenated sequence of tokens too, for multi-objective regression:

reg_lm = rlm.RegressLM.from_scratch(max_num_objs=2)

# Examples can have variable objective lengths.
examples = [core.Example(x='hello', y=[0.2]), core.Example(x='world', y=[-0.2, 0.3])]
reg_lm.fine_tune(examples)

# Now `samples` has shape (128, 2).
samples = reg_lm.sample([core.ExampleInput(x='hi')], num_samples=128)[0]

Pretrained Third-Party Models

T5Gemma (V1 + V2) encoder + our custom decoder is supported:

config = PyTorchModelConfig(architecture_kwargs=dict(encoder_type=EncoderType.T5GEMMA))

End-to-end T5Gemma (encoder + decoder) is also supported as a baseline:

from regress_lm.pytorch import t5gemma_model
model = t5gemma_model.T5GemmaModelConfig('google/t5gemma-s-s-prefixlm').make_model()

Long-Context

To support 100K+ input token lengths, alternative encoders (e.g. `mamba-ssm` and Performer) are supported:

architecture_kwargs = dict(encoder_type=EncoderType.MAMBA, additional_encoder_kwargs={'d_state': 128})
architecture_kwargs = dict(encoder_type=EncoderType.PERFORMER, additional_encoder_kwargs={'num_features': 256})

Citation

If you find this project useful, please consider citing the relevant works:

@article{code_rlm,
title={Regression Language Models for Code},
author={Yash Akhauri and Xingyou Song and Arissa Wongpanich and Bryan Lewandowski and Mohamed S. Abdelfattah},
booktitle={International Conference on Machine Learning},
series={Proceedings of Machine Learning Research},
year={2026},
}

@article{performance_prediction,
title={Performance Prediction for Large Systems via Text-to-Text Regression},
author={Yash Akhauri and Bryan Lewandowski and Cheng-Hsi Lin and Adrian N. Reyes and Grant C. Forbes and Arissa Wongpanich and Bangding Yang and Mohamed S. Abdelfattah and Sagi Perel and Xingyou Song},
journal={arXiv preprint arXiv:2506.21718},
year={2025}
}

@article{omnipred,
title={OmniPred: Language Models as Universal Regressors},
author={Xingyou Song and Oscar Li and Chansoo Lee and Bangding Yang and Daiyi Peng and Sagi Perel and Yutian Chen},
journal={Trans. Mach. Learn. Res.},
year={2024},
url={https://openreview.net/forum?id=t9c3pfrR1X},
}

@article{decoding_regression,
title={Decoding-based Regression},
author={Xingyou Song and Dara Bahri},
journal={Trans. Mach. Learn. Res.},
year={2025},
url={https://openreview.net/forum?id=avUQ8jguxg},
}

Disclaimer: This is not an officially supported Google product.