# How to Configure LogitsConfig to Return Embeddings for Downstream ML Tasks

> Easily configure LogitsConfig to return embeddings for downstream ML tasks by setting sequence and return_embeddings to True in the Biohub esm client. Logits() method extracts dense per-residue embeddings.

- Repository: [Biohub/esm](https://github.com/Biohub/esm)
- Tags: how-to-guide
- Published: 2026-05-30

---

**Set `sequence=True` and `return_embeddings=True` in the `LogitsConfig` dataclass to extract dense per-residue embeddings from ESM models via the `client.logits()` method.**

The Biohub/esm repository provides the `LogitsConfig` dataclass in [`esm/sdk/api.py`](https://github.com/Biohub/esm/blob/main/esm/sdk/api.py) to control which representations are computed during inference. When you configure LogitsConfig to return embeddings for downstream ML tasks, you enable the extraction of high-dimensional protein representations suitable for classification, clustering, or transfer learning without requiring separate model forward passes.

## Core LogitsConfig Fields for Embedding Extraction

Located in [[`esm/sdk/api.py`](https://github.com/Biohub/esm/blob/main/esm/sdk/api.py)](https://github.com/Biohub/esm/blob/main/esm/sdk/api.py#L14-L38), the `LogitsConfig` dataclass contains several boolean flags that determine which tensors are populated in the returned `LogitsOutput`:

- **`sequence`** – Must be set to `True` to request per-residue logits for the amino-acid sequence. This is required when extracting per-residue embeddings.
- **`return_embeddings`** – Returns the dense per-residue embeddings tensor (shape `[L, D]` where L is sequence length and D is embedding dimension). Set this to `True` to obtain embeddings for downstream classifiers or regressors.
- **`return_mean_embedding`** – Returns a single mean-pooled embedding vector (shape `[D]`) averaged over all residues. This is ideal for protein-level prediction tasks.
- **`return_hidden_states`** – Returns hidden states from each transformer layer. Useful when you need intermediate representations for analysis or layer-specific fine-tuning.
- **`ith_hidden_layer`** – Specific layer index to expose when using `return_hidden_states` (use `-1` for all layers, or zero-based index for specific layers).
- **`return_mean_hidden_states`** – Returns mean-pooled hidden states from the selected layer(s).

> **Architectural Note:** The `logits()` method in clients like `ESM3InferenceClient` or `ESMCInferenceClient` (defined in [`esm/sdk/forge.py`](https://github.com/Biohub/esm/blob/main/esm/sdk/forge.py) and [`esm/sdk/base_forge_client.py`](https://github.com/Biohub/esm/blob/main/esm/sdk/base_forge_client.py)) builds a `LogitsOutput` dataclass during the forward pass. When any `return_*` flag is enabled, the implementation populates the corresponding tensor fields without requiring additional model runs, making the same forward pass serve both traditional language modeling and representation learning.

## Extracting Per-Residue Embeddings

To obtain embeddings for each amino acid position in a protein sequence, enable both the sequence track and the embeddings flag:

```python
from esm.sdk.api import LogitsConfig, LogitsOutput
from esm.sdk.forge import client

# Encode a protein to tensor format

protein_tensor = client.encode(protein)

# Configure to return per-residue embeddings

config = LogitsConfig(sequence=True, return_embeddings=True)
output: LogitsOutput = client.logits(protein_tensor, config=config)

# embeddings tensor has shape [L, D]

per_residue_embeddings = output.embeddings

```

This pattern is validated in the test suite at [[`tests/oss_pytests/test_oss_client.py`](https://github.com/Biohub/esm/blob/main/tests/oss_pytests/test_oss_client.py)](https://github.com/Biohub/esm/blob/main/tests/oss_pytests/test_oss_client.py#L38-L40) and demonstrated in the cookbook example [[`cookbook/snippets/esm3.py`](https://github.com/Biohub/esm/blob/main/cookbook/snippets/esm3.py)](https://github.com/Biohub/esm/blob/main/cookbook/snippets/esm3.py#L6-L8).

## Generating Mean-Pooled Protein Embeddings

For downstream tasks requiring a single fixed-size vector per protein (such as whole-protein classification or clustering), configure the client to return mean-pooled embeddings:

```python
config = LogitsConfig(
    sequence=True,
    return_embeddings=True,
    return_mean_embedding=True,
)
output = client.logits(protein_tensor, config=config)

# Mean embedding is a [D] dimensional vector

protein_vector = output.mean_embedding

```

This approach aggregates residue-level information into a compact representation suitable for scikit-learn classifiers or deep learning architectures that expect fixed-length inputs.

## Extracting Intermediate Hidden States

When you need representations from specific transformer layers rather than the final output embeddings, use the hidden states configuration:

```python
config = LogitsConfig(
    sequence=True,
    return_hidden_states=True,
    ith_hidden_layer=5,  # Zero-based index for layer 5

    return_mean_hidden_states=True,
)
output = client.logits(protein_tensor, config=config)

# Layer-specific hidden states [L, D]

layer_activations = output.hidden_states

# Mean-pooled version [D]

layer_mean = output.mean_hidden_state

```

This technique is particularly valuable for probing tasks or when implementing layer-wise feature extraction for transfer learning pipelines.

## Summary

- **`LogitsConfig`** in [`esm/sdk/api.py`](https://github.com/Biohub/esm/blob/main/esm/sdk/api.py) controls which embeddings and hidden states are returned during inference.
- Set **`sequence=True`** and **`return_embeddings=True`** to extract per-residue embeddings for sequence-level downstream tasks.
- Enable **`return_mean_embedding`** to obtain a single pooled vector per protein suitable for classification and clustering.
- Use **`return_hidden_states`** with **`ith_hidden_layer`** to access intermediate transformer representations for layer-specific analysis.
- The `logits()` method populates **`LogitsOutput`** fields based on these flags without requiring multiple forward passes.

## Frequently Asked Questions

### What is the difference between `return_embeddings` and `return_hidden_states`?

**`return_embeddings`** returns the final output representations from the model's last layer, typically used as the default protein embedding. **`return_hidden_states`** returns intermediate activations from specified transformer layers, which are useful for analyzing how information propagates through the network or for layer-specific fine-tuning strategies.

### Do I need to set `sequence=True` to get embeddings?

**Yes.** The `sequence` field must be set to `True` because the embedding extraction logic is tied to the sequence track computation. Without enabling the sequence track, the model does not compute the representations necessary for populating the `embeddings` or `hidden_states` fields in the output.

### Can I use LogitsConfig with both ESM3 and ESMC models?

**Yes.** The `LogitsConfig` dataclass is designed to work across different ESM model implementations including `ESM3InferenceClient` and `ESMCInferenceClient`. However, note that the `sae_config` field is specific to ESM-C models for Sparse Auto-Encoder configurations and is not required for standard embedding extraction.

### What shape should I expect for the returned embedding tensors?

**Per-residue embeddings** return a tensor of shape `[L, D]` where $L$ is the protein sequence length and $D$ is the model's hidden dimension (e.g., 512, 1024, or 1536 depending on the ESM variant). **Mean embeddings** return a vector of shape `[D]`. **Hidden states** return `[L, D]` for specific layers or a stacked tensor when requesting all layers.