Skip to content

Pooling::Cls should use first real token (attention mask), not index 0 #213

Description

@amgix-dev

Problem

Pooling::Cls in rust/src/embeddings/local/pooling.rs always takes the hidden state at sequence index 0:

fn cls(output: &ModelOutput) -> Result<PooledOutputType, anyhow::Error> {
    match output {
        ModelOutput::Tensor(tensor) => tensor
            .get_on_dim(1, 0)  // always position 0
            ...
    }
}

For padded batches, CLS pooling should use the first non-padding token (via attention_mask), not a fixed index 0. This matters when sequences are left-padded.

Pooling::LastToken already has padding awareness in the same file; CLS does not.

Expected behavior

CLS pooling should:

  1. Take attention_mask as input (like mean pooling)
  2. Use the index of the first real token per sequence
  3. Match the behavior used by common ModernBERT embedding checkpoints (e.g. Granite, gte-modernbert)

Impact

CLS-pooled models may produce incorrect embeddings when batching with padding, especially left padding.

Suggested fix

Update Pooling::Cls to accept attention_mask and gather the first real token per row.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions