Extracting the outputs of highest probabilities from a model

[lenianiva]

In language modeling, a common diagnostics method is to list the top-10 or top-100 words with the highest probabilities, e.g. if a model tries to predict the next token in "What is ..."

f(["What", "is"]) -> ["this", "that", "it", ...]

Is there a way to extract say the top 10 outputs from a Turing model? e.g.

import Turing

@Turing.model function sequence()
    p ~ Turing.Beta(1, 1)
    if p < 0.3
        y ~ Turing.Categorical([0.9, 0.1])
        return (:a, y)
    else
        z ~ Turing.Categorical([0.1, 0.7, 0.2])
        return (:b, z)
    end
end

In this case the top 2 outputs are (:b, 1), (:a, 0). This is obviously not possible for continuous distributions.

In my use case there could be thousands of categories so listing all the possible outputs is infeasible.

[sunxd3]

In general, Turing won't work this out analytically (symbolically), but maybe you can trying sampling.

For instance,

julia> using Turing

julia> @model function sequence()
           p ~ Turing.Beta(1, 1)
           if p < 0.3
               y ~ Turing.Categorical([0.9, 0.1])
               return (:a, y)
           else
               z ~ Turing.Categorical([0.1, 0.7, 0.2])
               return (:b, z)
           end
       end

julia> model = sequence()

julia> chain = sample(model, Prior(), 1_000_000) # draw a million samples by simulating
Chains MCMC chain (1000000×4×1 reshape(::Matrix{Union{Missing, Float64}}, 1000000, 4, 1) with eltype Union{Missing, Float64}):

Iterations        = 1:1:1000000
Number of chains  = 1
Samples per chain = 1000000
Wall duration     = 3.35 seconds
Compute duration  = 3.35 seconds
parameters        = p, z, y
internals         = lp

Summary Statistics
  parameters      mean       std       mcse      ess_bulk      ess_tail       rhat   ess_per_sec 
      Symbol   Float64   Float64   Float64?      Float64?      Float64?   Float64?      Float64? 

           p    0.5003    0.2888     0.0003   997105.4129   999684.5459     1.0000   297288.4356
           z    2.1001    0.5384    missing       missing       missing    missing       missing
           y    1.1003    0.3004    missing       missing       missing    missing       missing

Quantiles
  parameters      2.5%     25.0%     50.0%     75.0%     97.5% 
      Symbol   Float64   Float64   Float64   Float64   Float64 

           p    0.0251    0.2501    0.5003    0.7508    0.9753
           z    1.0000    2.0000    2.0000    2.0000    3.0000
           y    1.0000    1.0000    1.0000    1.0000    2.0000

Because only one of y and z will appear with a particular run, there will be missing in the final chain.

Alternatively, a more direct way to is to

julia> outputs = []
Any[]

julia> for i in 1:1_000_000
           push!(outputs, model())
       end

julia> Turing.StatsBase.countmap(outputs)
Dict{Any, Int64} with 5 entries:
  (:b, 2) => 490338
  (:b, 3) => 139727
  (:a, 1) => 269973
  (:b, 1) => 70003
  (:a, 2) => 29959

BTW, what kind of application are you thinking?

[lenianiva]

My application for this is language modeling as probabilistic programs. In NLP, it is common to show the most probable next tokens.

[sunxd3]

Sounds interesting! Do you need to do posterior inference (the above code is only sampling from the prior distribution)?

[lenianiva]

Sounds interesting! Do you need to do posterior inference (the above code is only sampling from the prior distribution)?

No. For this use case I just need to sample from the prior.