Trace frame-level scores using lattice

[zrsjta]

Hi, team!

I have encountered a problem with k2 in my code. Below is the description of this problem.

For a nnet_output with shape [B, T, D], I am trying to calculate the scores on a graph (MMI numerator or denominator) with any prefix segment of nnet_output, namely nnet_output[:, :t, :], where t is any index smaller than T (the total length in time axis).

Currently, I implement it by a loop. But this leads to much computation. My code is below

        graph, _ = self.graph_compiler.compile(texts, self.P, replicate_den=True)
        T = x.size()[1]
        tot_scores = []
        for t in range(T, 0, -1):
            supervision = torch.Tensor([[0, 0, t]]).to(torch.int32) # [idx, start, length]
            dense_fsa_vec = k2.DenseFsaVec(nnet_output, supervision)
            lats = k2.intersect_dense(graph, dense_fsa_vec, output_beam=10.0)
            frame_score = lats.get_tot_scores(log_semiring=True, use_double_scores=True)
            tot_scores.append(frame_score)
        tot_scores = torch.cat(tot_scores)

Could these scores be calculated by parsing the lats obtained from the whole nnet_output, which means we can calculate them with only one k2.intersect_dense? Approximation is also ok for me.

Thanks for your help ! :)

[zrsjta]

Additional information:

1. we do NOT need this process differentiable.
2. It seems that the scores of each state and arc are accessible. Could we solve this problem by those scores on lattices?

[csukuangfj]

Please refer to the help doc of k2.intersect_dense:
https://k2-fsa.github.io/k2/python_api/api.html#k2.intersect_dense

There are two extra optional arguments:

def intersect_dense(a_fsas: Fsa,
                    b_fsas: DenseFsaVec,
                    output_beam: float,
                    a_to_b_map: Optional[torch.Tensor] = None,
                    seqframe_idx_name: Optional[str] = None,
                    frame_idx_name: Optional[str] = None) -> Fsa:

You can use

lats = k2.intersect_dense(graph, dense_fsa_vec, output_beam=10.0, seqframe_idx_name='seqframe', frame_idx_name='frame')

After the above call, the resulting lats has two extra 1-D tensor attributes: seqframe and frame.
That is, you can access them using

lats.seqframe
lats.frame

In your case, the supervision contains only one utterance, so lats.seqframe and lats.frame should be the same.
lats.frame contains values in the range from 0 to T and it contains as many entries as lats.num_arcs.

You can invoke k2.intersect_dense only once by feeding T frames and then call lats.get_forward_scores
to get the forward scores of each state. After this, you can use lats.frame to identify which states corresponding
to the t-th frame, and then use log_sum_exp to sum the scores of these states.

[EDITED]: The whole process is also differentiable.

[zrsjta]

Thanks for the reply

It seems that lats.frame represents the frame index of each arc. But I'm not sure we can recover the mapping from state to frame by lats.frame only.

Previously, I did not notice the API lats.frame. However, I have recovered this mapping in another way (see code below)

def trace_frame(lats):
    arcs = lats[0].as_dict()['arcs']

    frame2state = []
    prev_buf, cur_buf = [0], []

    for arc in arcs:
        f, t, _, _ = arc
        f, t = int(f), int(t)

        if f in prev_buf:
            if not t in cur_buf:
                cur_buf.append(t)

        else:
            frame2state.append(prev_buf)
            prev_buf = cur_buf
            cur_buf = [t]

    frame2state.append(prev_buf) # last frame
    frame2state.append([t]) # final state
    return frame2state

After that, I try to compute the frame-level score by frame2state, lats.get_forward_scores and then use log_sum_exp, but there is a large gap between the results obtained in this way and the way with a loop (as I presented first). Below is the code

        _, den = self.graph_compiler.compile(texts, self.P, replicate_den=True)

        T = x.size()[1]
        print(T)
        den_scores = []
        for t in range(T, 0, -1):
            supervision = torch.Tensor([[0, 0, t]]).to(torch.int32) # [idx, start, length]
            dense_fsa_vec = k2.DenseFsaVec(nnet_output, supervision)
            den_lats = k2.intersect_dense(den, dense_fsa_vec, output_beam=10.0)
            den_tot_scores = den_lats.get_tot_scores(log_semiring=True, use_double_scores=True)
            den_scores.append(den_tot_scores)
        den_scores = torch.cat(den_scores).unsqueeze(0) # [T] -> [B, T]
        print("den score computed from previous version: ", den_scores)

        # new implementation
        supervision = torch.Tensor([[0, 0, T]]).to(torch.int32)
        dense_fsa_vec = k2.DenseFsaVec(nnet_output, supervision)
        den_lats = k2.intersect_dense(den, dense_fsa_vec, output_beam=10.0,\
                   seqframe_idx_name='seqframe', frame_idx_name='frame')

        frame2state = trace_frame(den_lats)
        den_forward_scores = den_lats.get_forward_scores(log_semiring=True, use_double_scores=True)
        assert len(frame2state) == T + 2 # extra start and end state
        den_scores_new = []
        for t in range(T, 0, -1):
            states = frame2state[t]
            den_score = torch.logsumexp(den_forward_scores[states], dim=0)
            den_scores_new.append(den_score)
        den_scores_new = torch.stack(den_scores_new, dim=-1)
        print("den score computed from new version: ", den_scores_new)
        print("diffence of the two version: ", den_scores_new - den_scores)

And the results of the two versions cannot match.

den score computed from previous version:  tensor([[783.3993, 781.7219, 780.0333, 778.3439, 776.5803, 774.7105, 772.7158,
         770.5528, 768.1100, 762.4893, 754.7248, 746.1509, 736.5164, 726.7509,
         716.2331, 706.6452, 697.4263, 686.9293, 681.2793, 674.0765, 664.6529,
         657.2257, 650.4535, 643.3796, 634.3289, 628.0670, 619.4294, 609.9516,
         600.0689, 594.3592, 590.8794, 585.2605, 581.2088, 578.4872, 572.3086,
         571.1170, 566.4870, 558.6632, 552.9537, 545.4799, 541.9318, 530.6125,
         521.3177, 511.8524, 502.7617, 492.8600, 487.0916, 479.4112, 471.2255,
         462.2586, 454.3531, 442.9971, 436.3475, 428.7141, 420.9272, 413.6736,
         406.7801, 397.0022, 391.5931, 386.9415, 379.1681, 370.1025, 361.0875,
         351.7468, 341.6609, 333.3474, 325.8754, 318.5557, 309.9469, 303.1736,
         296.7173, 286.5587, 280.5947, 273.7012, 266.3182, 258.1839, 250.6710,
         242.8872, 236.6474, 229.7769, 221.9658, 213.3599, 204.9579, 192.9310,
         186.8552, 179.3116, 170.0203, 160.0337, 150.8213, 142.6257, 134.6330,
         127.6783, 116.1433, 103.3247,  91.5196,  80.5909,  69.7985,  59.0771,
          48.1386,  36.4054,  23.1133,  11.1909,  -2.9392]],
       dtype=torch.float64)
den score computed from new version:  tensor([783.3993, 781.7252, 780.0382, 778.3510, 776.5909, 774.7267, 772.7416,
        770.5969, 768.1950, 762.6968, 755.2699, 747.3608, 738.8629, 730.4765,
        725.2539, 716.5815, 707.1374, 698.3506, 691.5351, 679.1446, 670.6063,
        663.6814, 656.9236, 651.3605, 647.0456, 640.6885, 626.9087, 617.5355,
        610.5875, 604.0881, 598.1472, 590.7349, 588.7432, 585.2757, 583.0598,
        581.2615, 570.5313, 566.4218, 561.7389, 557.9919, 554.6887, 543.3505,
        529.1934, 521.2846, 512.1856, 504.0500, 500.0110, 490.2833, 477.4947,
        471.4613, 463.3775, 455.1872, 449.6866, 438.6009, 426.2327, 421.9648,
        415.6079, 408.7287, 403.9518, 395.7307, 387.3507, 378.8256, 370.0204,
        361.4262, 355.0736, 346.0511, 331.8522, 322.9216, 317.4937, 311.1927,
        304.5424, 297.9715, 292.8139, 282.6298, 272.3836, 264.5364, 257.1948,
        250.8738, 245.0846, 235.5465, 227.8981, 221.6154, 213.5615, 205.2312,
        199.0696, 187.7953, 177.1561, 169.5724, 160.6012, 151.7466, 146.6541,
        140.5165, 129.1156, 116.3533, 104.6344,  93.8204,  83.1367,  72.5044,
         61.6243,  49.9238,  36.6502,  24.7326,  10.5949], dtype=torch.float64)
diffence of the two version:  tensor([[    0.0000,     0.0034,     0.0049,     0.0071,     0.0106,     0.0162,
             0.0258,     0.0441,     0.0850,     0.2075,     0.5451,     1.2099,
             2.3464,     3.7256,     9.0209,     9.9363,     9.7111,    11.4214,
            10.2558,     5.0681,     5.9534,     6.4558,     6.4701,     7.9809,
            12.7167,    12.6216,     7.4792,     7.5838,    10.5185,     9.7289,
             7.2677,     5.4744,     7.5344,     6.7885,    10.7512,    10.1445,
             4.0443,     7.7586,     8.7853,    12.5120,    12.7569,    12.7380,
             7.8757,     9.4322,     9.4238,    11.1901,    12.9194,    10.8722,
             6.2692,     9.2026,     9.0244,    12.1901,    13.3391,     9.8868,
             5.3055,     8.2912,     8.8277,    11.7265,    12.3586,     8.7891,
             8.1826,     8.7230,     8.9329,     9.6794,    13.4127,    12.7037,
             5.9769,     4.3660,     7.5468,     8.0192,     7.8250,    11.4128,
            12.2193,     8.9286,     6.0653,     6.3525,     6.5239,     7.9867,
             8.4372,     5.7696,     5.9323,     8.2554,     8.6036,    12.3002,
            12.2144,     8.4837,     7.1358,     9.5387,     9.7798,     9.1209,
            12.0211,    12.8382,    12.9724,    13.0286,    13.1148,    13.2295,
            13.3382,    13.4273,    13.4856,    13.5184,    13.5370,    13.5417,
            13.5341]], dtype=torch.float64)

Currently, I don't know which version is correct.
The motivation of this is to compute the probability P(W|O_{1:t}) according to the definition of LF-MMI but only use the first several frames.

[csukuangfj]

arcs = lats.arcs.values()[:, :2]
# arcs is a 2-D torch.int32 tensor



for idx, (src, dst) in enumerate(arcs.tolist()):
  # note src is not used and you can replace it with an underscore _
   frame_idx = lats.frame[idx]
   # now you konw the state `dst` belongs to the frame `frame_idx`
  # You can add the forward_score of this state to a list corresponding to the frame `frame_idx`
  # 
  # Caution: You have to avoid adding `dst` state multiple times 

# At this point, you know the states corresponding to each frame, you can use `log-sum-exp` to combine them.

[csukuangfj]

But I'm not sure we can recover the mapping from state to frame by lats.frame only.

As I posted above, you can iterate over the arcs; for each arc, you can get its frame_idx and dest_state.
This dest_state belongs to the frame frame_idx.

If you have multiple utterances, then you have to use seqframe_idx.

[csukuangfj]

Note: I have re-edited the demo code.

[csukuangfj]

Previously, I did not notice the API lats.frame. However, I have recovered this mapping in another way (see code below)

Please use a small lats to verify that your code is correct.
(You can print the resulting lats, its frame, its states, the return value frame2state to check the correctness of your code)

Looks like your frame2state is a 1-d list, which is not correct.

A frame can correspond to multiple states, while a state belongs to only one frame.

[csukuangfj]

You can note down which state belongs to which frame.
For example

frame_0_states = [1, 2, 3]
frame_1_states = [4, 5, 6, 7, 8, 9]
....

You can get the total scores for frame 0 using

forward_scores = lats.get_forward_scores(use_double_scores=True, log_semiring=True)

frame_0_tot_scores = forward_scores[frame_0_states].exp().sum().log()

Note: For the last frame T-1, you have to consider the scores on the arcs entering the final state if those scores are not zero.

[zrsjta]

@csukuangfj
As suggested, I restrict the nnet_output to 3 frames and checked the frame2state. I suppose my frame2state is correct.

den score computed from previous version:  tensor([[23.1133, 11.1909, -2.9392]], dtype=torch.float64)
arcs: tensor([[          0,           1,           0,  1083890428],
        [          0,           2,         147,  1093236685],
        [          1,          13,         147,  1093883163],
        [          2,           3,          21, -1083650674],
        [          2,           4,          35,  1042883936],
        [          2,           5,          43, -1103412784],
        [          2,           6,          48, -1091424240],
        [          2,           7,         111, -1081929866],
        [          2,           8,         112, -1090075914],
        [          2,           9,         114,  1050298232],
        [          2,          10,         115, -1070879172],
        [          2,          11,         140, -1089980508],
        [          2,          12,         141, -1078741436],
        [          2,          13,         147,  1096957993],
        [          2,          14,         153, -1097704528],
        [          2,          15,         189, -1077760680],
        [          3,          16,           1,  1081681798],
        [          4,          16,           1,  1079942965],
        [          5,          16,           1,  1078807868],
        [          6,          16,           1,  1082282282],
        [          7,          16,           1,  1080307114],
        [          8,          16,           1,  1080573812],
        [          9,          16,           1,  1082305666],
        [         10,          16,           1,  1085392658],
        [         11,          16,           1,  1081433908],
        [         12,          16,           1,  1082813821],
        [         13,          18,           0,  1092706735],
        [         13,          16,           1, -1058648121],
        [         13,          17,         147,  1094443926],
        [         14,          16,           1,  1081785469],
        [         15,          16,           1,  1082840980],
        [         16,          19,          -1, -1221591040],
        [         17,          19,          -1, -1051153366],
        [         18,          19,          -1, -1051153366]],
       dtype=torch.int32)
lats.frame:  tensor([0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2, 2,
        2, 2, 2, 2, 2, 2, 2, 3, 3, 3], dtype=torch.int32)
frame2state:  [[0], [1, 2], [13, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 15], [16, 18, 17], [19]]
den score computed from new version:  tensor([36.6502, 24.7326, 10.5949], dtype=torch.float64)
diffence of the two version:  tensor([[13.5370, 13.5417, 13.5341]], dtype=torch.float64)

[zrsjta]

Note: For the last frame T-1, you have to consider the scores on the arcs entering the final state if those scores are not zero.

I'm also worried about the scores on the final states:
If nnet_output[:, :t, :] is used in intersect_dense, states that belong to t-th frame should be considered as the final states and final scores should be considered. (In k2 you implement it by an additional final state with these final scores on some additional arcs). This is what happens in my original implementation.
However, by tracing the lattice, states that belong to all frames (except the last frame) are not considered as final states and don't add the extra final scores. That may cause a difference in the scores.

[csukuangfj]

frame2state: [[0], [1, 2], [13, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 15], [16, 18, 17], [19]]

[0] is the start state and does not belong to any frame. Looks like there is an offset-by-one error.

For arc i, we get its dest_state (not src_state) and frame_idx = lats.frame[i], then dest_state belongs to frame frame_idx.

[zrsjta]

I have dumped the lattice as below.
if the frame is indexed from 1, i suppose we should do logsumexp on groups below.
frame_1_states = [1, 2]
frame_2_states = [13, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 15]
frame_3_states = [16, 18, 17]

[0] is the start state and does not belong to any frame. Looks like there is an offset-by-one error.

In my frame2state, the start and end states are the first and last entries respectively. Note the loop for t in range(T, 0, -1) with T=3 means only t=3, 2, 1 would be called, which would ignore both [0] and [19]

For arc i, we get its dest_state (not src_state) and frame_idx = lats.frame[i], then dest_state belongs to frame frame_idx.

Follow this advice, I have also revised my trace_frame function as below

def trace_frame(lats):
    arcs = lats.arcs.values()[:, :2]
    T = max(lats.frame).item()
    frame2state = [[] for _ in range(T+1)]

    for idx, (_, dst) in enumerate(arcs.tolist()):
        frame_idx = lats.frame[idx]
        if dst not in frame2state[frame_idx]:
            frame2state[frame_idx].append(dst)

    print("frame2state: ", frame2state)
    return frame2state

and the output of this is like below. It only ignore the start state [0] while other entries are the same
frame2state: [[1, 2], [13, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 15], [16, 18, 17], [19]]

Use the new frame2state, the difference between the two versions are unchanged:

den score computed from previous version:  tensor([[23.1133, 11.1909, -2.9392]], dtype=torch.float64)
frame2state:  [[1, 2], [13, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 15], [16, 18, 17], [19]]
logsumexp on [16, 18, 17] is 36.650240146578824
logsumexp on [13, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 15] is 24.732586652356957
logsumexp on [1, 2] is 10.594917989250451
den score computed from new version:  tensor([36.6502, 24.7326, 10.5949], dtype=torch.float64)
diffence of the two version:  tensor([[13.5370, 13.5417, 13.5341]], dtype=torch.float64)

Note:
I have changed assert len(frame2state) == T + 2 to assert len(frame2state) == T + 1 since [0] does not exits now
I have changed states = frame2state[t] to states = frame2state[t-1] also because [0] does not exits now
The log above presents all states in each logsumexp. Now I suppose we don't have the ``offset-by-one error

[zrsjta]

Note: For the last frame T-1, you have to consider the scores on the arcs entering the final state if those scores are not zero.

As shown in the lattice above, arc 17->19 and 18->19 have scores -13.54. So the scores on state 17 and 19 shoude be revised before logsumexp. (So far i don't do it)
I'm curious whether a similar modification should be done for states that belong to other frames. As in my privous implementation, each frame would be considered as the final frame once. If needed, how to do this?

Thanks :)

[csukuangfj]

I just created a colab notebook (see https://colab.research.google.com/drive/1iyc_q8aHuKd-RZxtYv9EqfyjB2QZDSOx?usp=sharing)
to verify the idea.

The following code you posted:

        graph, _ = self.graph_compiler.compile(texts, self.P, replicate_den=True)
        T = x.size()[1]
        tot_scores = []
        for t in range(T, 0, -1):
            supervision = torch.Tensor([[0, 0, t]]).to(torch.int32) # [idx, start, length]
            dense_fsa_vec = k2.DenseFsaVec(nnet_output, supervision)
            lats = k2.intersect_dense(graph, dense_fsa_vec, output_beam=10.0)
            frame_score = lats.get_tot_scores(log_semiring=True, use_double_scores=True)
            tot_scores.append(frame_score)
        tot_scores = torch.cat(tot_scores)

is not equivalent to the one where the tot_scores is computed from the whole lattice (when you feed T frames at once)

[zrsjta]

I just created a colab notebook (see https://colab.research.google.com/drive/1iyc_q8aHuKd-RZxtYv9EqfyjB2QZDSOx?usp=sharing)
to verify the idea.

The following code you posted:

        graph, _ = self.graph_compiler.compile(texts, self.P, replicate_den=True)
        T = x.size()[1]
        tot_scores = []
        for t in range(T, 0, -1):
            supervision = torch.Tensor([[0, 0, t]]).to(torch.int32) # [idx, start, length]
            dense_fsa_vec = k2.DenseFsaVec(nnet_output, supervision)
            lats = k2.intersect_dense(graph, dense_fsa_vec, output_beam=10.0)
            frame_score = lats.get_tot_scores(log_semiring=True, use_double_scores=True)
            tot_scores.append(frame_score)
        tot_scores = torch.cat(tot_scores)

is not equivalent to the one where the tot_scores is computed from the whole lattice (when you feed T frames at once)

My observation is the same: the tot_scores computed by the two methods are different.
Is there anything wrong with my code? or any suggestions on this mismatch?