Embedding example - confusion about normalization

[MartinPerry]

In the official embedding example, there si this loop:

    for (int i = 0; i < batch.n_tokens; i++) {
        if (!batch.logits[i]) {
            continue;
        }

        const float * embd = nullptr;
        int embd_pos = 0;

        if (pooling_type == LLAMA_POOLING_TYPE_NONE) {
            // try to get token embeddings
            embd = llama_get_embeddings_ith(ctx, i);
            embd_pos = i;
            GGML_ASSERT(embd != NULL && "failed to get token embeddings");
        } else {
            // try to get sequence embeddings - supported only when pooling_type is not NONE
            embd = llama_get_embeddings_seq(ctx, batch.seq_id[i][0]);
            embd_pos = batch.seq_id[i][0];
            GGML_ASSERT(embd != NULL && "failed to get sequence embeddings");
        }

        float * out = output + embd_pos * n_embd;
        common_embd_normalize(embd, out, n_embd, embd_norm);
    }

It runs over batch.n_tokens, which is number of tokens in the input data. But why is common_embd_normalize called for every token?

[ngxson]

There is a if (!batch.logits[i]) condition on the top of this loop. The way it works is that for all the tokens in sentence, there is no logits, so we skip. Except for the last token where the logits is actually the embedding output of all past tokens + the last token.

[MartinPerry]

Hm.. but I have all logits set to 1. In

static void batch_add_seq(llama_batch & batch, const std::vector<int32_t> & tokens, llama_seq_id seq_id) {
    size_t n_tokens = tokens.size();
    for (size_t i = 0; i < n_tokens; i++) {
        common_batch_add(batch, tokens[i], i, { seq_id }, true);
    }
}

There is true, which results in batch.logits[batch.n_tokens] = logits;. If it should be rewritten by llama_decode, it is not and I have the same array.

[ngxson]

Hmm ok that seems to be a bug (or I have been misunderstand about embeddings API all along?). Do you have any ideas @ggerganov ?

[ggerganov]

Looks like a bug - true. When pooling is used, it will output the same thing for all the tokens, instead of just once for each sequence.

I think we fixed this in the upcoming #11875?

[ngxson]

Yes I pushed a fix: 8ec0ff9