You should be able to calculate the entropy directly from log probs as

entropy = -(per_token_logps.exp() * per_token_logps).sum(-1)

which means we don't have to modify _get_per_token_logps. I noted however that you didn't apply temperature scaling when calculating your entropy, which my proposed solution does. As I see it temperature scaling remains an unresolved variable: the paper evaluates entropy masking only at T = 1, so the optimal percentile thresholds for other temperatures are unknown. The same caveat applies to model choice—results are reported solely for the Qwen-3 series, and the threshold may shift substantially for other architectures. In practice, the 20 % rule should be treated as a starting heuristic; users will need to experiment with the temperature-model-percentile combination that best fits their own setup.

Ah, yeah that's a good point let me update the entropy calculation to use the logprobs directly. Yep the temperature and 20% threshold should be hyper-parameters that are tuned by the user.

Also just realized that the temperature parameter won't affect the ranking order of entropy since all positions are affected by the same temperature, so the temp param shouldn't matter here.

You should be able to calculate the entropy directly from log probs as

entropy = -(per_token_logps.exp() * per_token_logps).sum(-1)

Actually this isn't true the per_token_logps only contain the log probs of the token in the sampled response. Entropy needs to be calculated over the probability distribution over the entire vocabulary per seq position in the response, so we do need to update the _get_per_token_logps since it uses the selective_log_softmax function to gather the probs of just the tokens in sampled response.

You should be able to calculate the entropy directly from log probs as

entropy = -(per_token_logps.exp() * per_token_logps).sum(-1)

Actually this isn't true the per_token_logps only contain the log probs of the token in the sampled response. Entropy needs to be calculated over the probability distribution over the entire vocabulary per seq position in the response, so we do need to update the _get_per_token_logps since it uses the selective_log_softmax function to gather the probs of just the tokens in sampled response.

yes of course, my bad. not really a fan of the proposed refactor of _get_per_token_logps so trying to think of an alternative solution. Things of concern: 1) logits is a big tensor, clogs VRAM. 2) performing torch.softmax(logits, dim=-1) will spike memory usage. See #3390 for discussion. But maybe we just have to bite the bullet here, and in this case we should definitely only perform the entropy calculation when necessary

Nice! Thanks!
Yes, as @LeonEricsson says, the tricky part is the memory peak of the softmax. There as to be a way to compute the entropy efficiently, similarly to what we do with selective_log_softmax.

Another recommendation: 1 argument is probably enough.

- if self.filter_on_entropy:
+ if self.token_entropy_percentile_threshold < 1.0:

and add that the recommended value is 0.2 in the documentation.

Cool, let me look into making the entropy calculation less memory intensive.

Updated the code to make sure that only a mini-batch of logits are materialized at any given point of time and entropies for those mini-batches of logits are optionally calculated. The rowise_entropy function is inspired by the selective_log_softmax function to calculate entropy per row of the mini-batch. Introduces a bit of redundant compute because we compute the log_softmax twice in entropy calc and selective_log_softmax, but I guess its worth the tradeoff with the gains in clean code and memory.

Updated the code to make sure that only a mini-batch of logits are materialized at any given point of time and entropies for those mini-batches of logits are optionally calculated. The rowise_entropy function is inspired by the selective_log_softmax function to calculate entropy per row of the mini-batch. Introduces a bit of redundant compute because we compute the log_softmax twice in entropy calc and selective_log_softmax, but I guess its worth the tradeoff with the gains in clean code and memory.

nice work. left a few minor comments

Thanks for the review Leon! Made the suggested changes.

@qgallouedec please take another look at the PR when you have the time.

I spent some time benchmarking different entropy calculations, scripts here. Long story short I recommend:

def compute_entropy(logits: torch.Tensor, chunk_size: int = 1) -> torch.Tensor:
    """
    Compute the Shannon entropy (nats) for each row of `logits`
    with a memory-friendly, `log_softmax → exp → reduce` sequence.

    Args:
        logits (`torch.Tensor`):
            Tensor of shape `(..., num_classes)` holding unnormalised
            log-probabilities.
        chunk_size (`int`, *optional*, defaults to `1`):
            Number of rows processed together. Keeping it
            at `1` minimises peak memory; larger values trade a small
            rise in memory for fewer launches.  Default is `1`.

    Returns:
        `torch.Tensor`:
            Entropy values with shape `logits.shape[:-1]`.
    """
    outs = []
    for chunk in logits.split(chunk_size, dim=0): # loop to reduce peak mem consumption
        logp = F.log_softmax(chunk, dim=-1)
        entropy = -(logp.exp() * logp).sum(-1)
        outs.append(entropy)
    return torch.cat(outs, dim=0)

this is basically what you already had, but importantly we sum inside the loop which avoids materializing the [S, V] tensor. I observed no latency improvements when increasing chunk_size on a 4090, because it's already memory-bound at 1

@qgallouedec @LeonEricsson this PR has a bit of duplicate with my implementation of entropy regularization loss in #3628. We need to sync on both given the 2 address entropy in two different directions. The entropy regularization loss is proposed in issue and also officially implemented in verl.

@1485840691 I think that despite the overlap both of the purposes are complementary and once this PR is pushed including the entropy loss in the final loss should be fairly simple.

@LeonEricsson thanks for those cool benchmarks! I've updated the code to compute the sum inside the for loop as you've suggested.

I have a comment on entropy from logits. Given verl has already provided an implementation https://github.com/volcengine/verl/blob/9b7bb69ea3165b691cc908d7f3f2f14c4a65a59e/verl/utils/torch_functional.py#L150, why do we not re-use that?

Sorry, I tried benchmark it. Indeed verl's implementation does not have better running performance
"
Duration of rowise_entropy: 0.7136 seconds
Duration of entropy_from_logits: 1.4799 seconds
"

But given the entropy from logits function is embed inside the chunk loop of get_per_token_logps, why do we need to support chunking here?

And there is another question regarding the interaction between entropy loss and entropy mask: Do we need to consider the entropy mask in computing entropy loss? Now entropy loss is computed using completion mask @qgallouedec @LeonEricsson

I have a comment on entropy from logits. Given verl has already provided an implementation https://github.com/volcengine/verl/blob/9b7bb69ea3165b691cc908d7f3f2f14c4a65a59e/verl/utils/torch_functional.py#L150, why do we not re-use that? Given the entropy from logits function is embed inside the chunk loop of get_per_token_logps, I think we might not need to support chunking here. Even if we do need chunking, verl has in-place implementation: https://github.com/volcengine/verl/blob/9b7bb69ea3165b691cc908d7f3f2f14c4a65a59e/verl/utils/torch_functional.py#L150

I'm assuming you mean why we don't simply import the function from veRL. I think that we'd like to avoid introducing new dependencies for fairly simple utility functions. Good point on supporting chunking inside the entropy function, while you correctly point out that we don't need to support chunking for the current implementation of GRPO since the get_per_token_logps does it, I think it makes sense to support chunking inside the rowise_entropy function to make it more re-usable for other usecases.

And there is another question regarding the interaction between entropy loss and entropy mask: Do we need to consider the entropy mask in computing entropy loss? Now entropy loss is computed using completion mask @qgallouedec @LeonEricsson

I don't think the entropy loss should take the entropy mask into consideration. In the entropy masking paper linked above, their intention was to just affect the policy loss. So similar to how KL-div doesn't consider the entropy mask, I believe that the entropy loss shouldn't either, especially if we want to reproduce the behavior of veRL.

I have a comment on entropy from logits. Given verl has already provided an implementation https://github.com/volcengine/verl/blob/9b7bb69ea3165b691cc908d7f3f2f14c4a65a59e/verl/utils/torch_functional.py#L150, why do we not re-use that? Given the entropy from logits function is embed inside the chunk loop of get_per_token_logps, I think we might not need to support chunking here. Even if we do need chunking, verl has in-place implementation: https://github.com/volcengine/verl/blob/9b7bb69ea3165b691cc908d7f3f2f14c4a65a59e/verl/utils/torch_functional.py#L150

I'm assuming you mean why we don't simply import the function from veRL. I think that we'd like to avoid introducing new dependencies for fairly simple utility functions. Good point on supporting chunking inside the entropy function, while you correctly point out that we don't need to support chunking for the current implementation of GRPO since the get_per_token_logps does it, I think it makes sense to support chunking inside the rowise_entropy function to make it more re-usable for other usecases.

I think we might follow verl to support entropy_from_logits(https://github.com/volcengine/verl/blob/9b7bb69ea3165b691cc908d7f3f2f14c4a65a59e/verl/utils/torch_functional.py#L143) and entropy_from_logits_chunking(https://github.com/volcengine/verl/blob/9b7bb69ea3165b691cc908d7f3f2f14c4a65a59e/verl/utils/torch_functional.py#L150). And I do not mean importing verl lib to support such small util function ,just copy the code and add a comment. But given I have benchmarked it. verl's implementation does not have better performance.

I have a comment on entropy from logits. Given verl has already provided an implementation https://github.com/volcengine/verl/blob/9b7bb69ea3165b691cc908d7f3f2f14c4a65a59e/verl/utils/torch_functional.py#L150, why do we not re-use that?

Sorry, I tried benchmark it. Indeed verl's implementation does not have better running performance " Duration of rowise_entropy: 0.7136 seconds Duration of entropy_from_logits: 1.4799 seconds "

But given the entropy from logits function is embed inside the chunk loop of get_per_token_logps, why do we need to support chunking here?

Sorry, I missed this comment. I think the objective here is to reduce peak memory usage rather than optimize on run-time, based on the comments from Leon and Quentin

Nice! Thanks!
Yes, as @LeonEricsson says, the tricky part is the memory peak of the softmax. There as to be a way to compute the entropy efficiently, similarly to what we do with selective_log_softmax.

But given the entropy from logits function is embed inside the chunk loop of get_per_token_logps, why do we need to support chunking here?

I agree that chunking isn't necessary here, but since it's defined as a general utility function and could be useful in future scenarios, its nice to have; provides a convenient way to trade memory for latency.

I think we might follow verl to support entropy_from_logits(https://github.com/volcengine/verl/blob/9b7bb69ea3165b691cc908d7f3f2f14c4a65a59e/verl/utils/torch_functional.py#L143) and entropy_from_logits_chunking(https://github.com/volcengine/verl/blob/9b7bb69ea3165b691cc908d7f3f2f14c4a65a59e/verl/utils/torch_functional.py#L150).

Yeah, entropy_from_logits_chunking looks very similar to what I benchmarked here. I’d be fine with it as well, but based on my isolated benchmarks, I found it has the same memory usage while being slower.

I don't think the entropy loss should take the entropy mask into consideration. In the entropy masking paper linked above, their intention was to just affect the policy loss. So similar to how KL-div doesn't consider the entropy mask, I believe that the entropy loss shouldn't either, especially if we want to reproduce the behavior of veRL.

looking into this
I agree with @pramodith, entropy filtering is applied to the policy loss, and the entropy loss should be entirely separate from the policy loss.

some final comments then I'm happy to merge.

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.