dynamic-frame-compression/writeup.html at main · floatingtrees/dynamic-frame-compression · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="utf-8">
<title>Dynamic Frame Compression</title>
<script>
window.MathJax = {
  tex: { inlineMath: [['$', '$'], ['\\(', '\\)']], displayMath: [['$$','$$'], ['\\[','\\]']] },
  svg: { fontCache: 'global' }
};
</script>
<script id="MathJax-script" async src="https://cdn.jsdelivr.net/npm/mathjax@3/es5/tex-mml-chtml.js"></script>
<style>
  :root { --fg:#1a1a1a; --muted:#5a5a5a; --accent:#2c5aa0; --bg:#fdfdfd; --border:#e3e3e3; --code-bg:#f4f4f4; }
  html { font-size: 17px; }
  body {
    font-family: 'Charter', 'Iowan Old Style', 'Palatino Linotype', Georgia, serif;
    line-height: 1.65;
    max-width: 760px;
    margin: 3rem auto;
    padding: 0 1.25rem 4rem;
    color: var(--fg);
    background: var(--bg);
  }
  h1 { font-size: 2.1rem; margin: 0 0 0.4rem; line-height: 1.2; }
  h2 { font-size: 1.45rem; margin-top: 2.4rem; border-bottom: 1px solid var(--border); padding-bottom: 0.3rem; }
  h3 { font-size: 1.15rem; margin-top: 1.6rem; }
  p { margin: 0.9rem 0; }
  a { color: var(--accent); }
  code { font-family: 'JetBrains Mono', 'Menlo', Consolas, monospace; font-size: 0.88em;
         background: var(--code-bg); padding: 0.1em 0.35em; border-radius: 3px; }
  figure { margin: 1.4rem 0; text-align: center; }
  figure img, figure video { max-width: 100%; height: auto; display: block; margin: 0 auto;
                              border: 1px solid var(--border); border-radius: 4px; }
  figcaption { font-size: 0.88rem; color: var(--muted); margin-top: 0.4rem; font-style: italic; }
  table { border-collapse: collapse; margin: 1.2rem auto; font-size: 0.93rem; }
  th, td { border: 1px solid var(--border); padding: 0.4rem 0.7rem; text-align: left; }
  th { background: var(--code-bg); font-weight: 600; }
  td.num, th.num { text-align: right; font-variant-numeric: tabular-nums; }
  sup.fn { font-size: 0.75em; }
  sup.fn a { text-decoration: none; color: var(--accent); }
  .footnotes { margin-top: 3rem; padding-top: 1rem; border-top: 1px solid var(--border);
               font-size: 0.9rem; color: var(--muted); }
  .footnotes ol { padding-left: 1.4rem; }
  .footnotes li { margin: 0.5rem 0; }
  .footnotes li a.fn-back { text-decoration: none; margin-left: 0.3rem; }
  .meta { color: var(--muted); margin-bottom: 2rem; }
  .grid-row { display: flex; gap: 0.6rem; flex-wrap: wrap; justify-content: center; margin: 1rem 0; }
  .grid-row figure { flex: 1 1 320px; margin: 0; }
</style>
</head>
<body>

<h1>Dynamic Frame Compression</h1>
<p class="meta">A video autoencoder with a dynamic temporal compression rate.</p>

<p>Video generation models are expensive to run. To naively generate 1 minute of video at 32 fps and 720p resolution, we need to perform attention over a sequence of 1920 frames and 2.7 million features per frame, which takes 5 GB of memory per layer to store at fp8. If we use an aggressive image autoencoder such as <a href="https://arxiv.org/abs/2410.10733">DC-AE</a> with around a 25x compression rate, we can cut the number of features per frame to around 100k.</p>

<p>However, this still wastes a large amount of memory because neighboring frames contain a lot of uninteresting content. We can deal with this problem by training a video autoencoder that compresses both spatial and temporal dimensions.</p>

<p>Most video autoencoders use a fixed spatial and temporal compression ratio. This means that they compress a video in shape of <code>(height, width, number of frames, RGB_channels)</code> to a tensor in shape of <code>(height / spatial_compression_rate, width / spatial_compression_rate, number of frames / temporal_compression_rate, channels)</code>.<sup class="fn"><a id="fnref-1" href="#fn-1">1</a></sup> However, some videos are much more compressible than others; a slow zoom-in contains much less information than an action movie scene. This means that a large amount of compute is wasted on uninteresting computation.</p>

<p>To bypass this, the current mainstream approach is to use a model to pick out key frames of a video, train the diffusion model to only generate keyframes, and during inference, run a cheap video model to fill in the intermediate frames between each keyframe. Although this cuts down on computation, it reduces the signal given to the diffusion model and prevents the diffusion model from controlling the entire scene, leading to slow and predictable movements.</p>

<p>Our solution is using a video autoencoder with a dynamic temporal compression rate. During compression, the autoencoder compresses the video into a latent space and selects some elements of the latent space to mask out. To be precise, our encoder maps the video from <code>(height, width, number of frames, RGB_channels)</code> to a tensor in shape of <code>(height / spatial_compression_rate, width / spatial_compression_rate, number of selected frames, channels)</code>, and a boolean mask in shape of <code>(total_frames, )</code>. Since the compression happens after a sequence of spatiotemporal attention layers, the encoder can move information between frames before masking, which allows the diffusion model to see the entire video.</p>

<h2>Implementation Details</h2>

<p>The primary difficulty was training the encoder to mask frames properly. Deleting frames is implemented as multiplying a compressed representation in shape of <code>(height / spatial_compression_rate, width / spatial_compression_rate, total_frames, channels)</code> with a boolean mask in shape of <code>(total_frames, )</code>. The mask is generated by Bernoulli sampling each element of the <code>(total_frames, )</code> tensor. Thus, if we train the autoencoder to generate a reasonable mask, everything should fall into place around it. To prevent the model from keeping every frame, we add a L1 penalty, $\lambda_1$ on the sum of the boolean mask, so $L = L_{\text{reconstruction}} + \lambda_1 \cdot \text{number\_of\_kept\_frames}$.<sup class="fn"><a id="fnref-2" href="#fn-2">2</a></sup></p>

<p>We started with a STE (straight through estimator). If we applied a Gumbel-sigmoid activation function on the logits of the mask, the model should learn which frames would minimize reconstruction loss the most. Unfortunately, this did not work; the model simply kept the same set of frames for every video. We suspect this failed because this problem is too discrete for the STE to handle; gumbel softmax is typically used in MoE training, where using a convex combination of multiple experts is acceptable. However, in our regime, the STE was only able to see gradients from entirely masked or unmasked frames, without any in-between frames.<sup class="fn"><a id="fnref-3" href="#fn-3">3</a></sup></p>

<p>Left without any other choice, we were forced to use RL to solve this problem. We trained a value network that predicts the expected loss of each video, and used that as the expectation baseline for RL. This worked, but it was unstable, added another model to train, and required frequent tuning when stacking with other changes.</p>

<p>Ultimately, we chose to simply generate two masks by sampling twice, compute loss for both, and use the mean of both masks as the expected reward, which is empirically much more stable. This doubles the compute requirements for the decoder, but the decoder is typically a larger model and requires more signal anyways. In our case, we would have required multiple epochs over the dataset, so a 1:2 encoder to decoder data ratio was deemed acceptable.</p>

<p>Regardless of which method we chose, the encoder had a strong tendency to create degenerate probability distributions, where the P(frame is masked) approached 0 or 1. This is likely the optimal solution; some frames are critical for reconstruction, and other frames are near useless. However, degenerate distributions prevent exploration and passing identical latents into the decoder wastes compute.</p>

<p>We speculated this happened because of the following phenomenon: if two frames were identical, the decoder would prefer to assign a probability of 1 to one of the frames and 0 to the other frame. This probability assignment guarantees exactly one frame is kept, so there is no risk of keeping zero frames (explodes reconstruction loss) or two frames (increases L1 penalty for no extra information). However, if two frames were equal, we might prefer a distribution where both frames have a 50% chance of being kept to preserve exploration.</p>

<p>An intuitive way to resolve this problem is with stratified sampling. Each sequence of mask probabilities was chunked into subarrays with total probability of 1. For example, if we had a vector of probabilities that looked like <code>[0.5, 0.5, 0.3, 0.4, 0.2, 0.1, 0.5, 0.1]</code>, the array would be split into <code>[0.5, 0.5]</code>, <code>[0.3, 0.4, 0.2, 0.1]</code>, and <code>[0.5, 0.1]</code>. We would sample an element from each subarray before merging them back; in this example, the sample might look like <code>[1, 0, 0, 0, 1, 0, 0, 0]</code>.<sup class="fn"><a id="fnref-4" href="#fn-4">4</a></sup></p>

<p>Stratified sampling can be efficiently implemented on modern hardware, but the number of unmasked frames is near constant, so the model faces significantly less pressure to learn high compression rates. Empirically, the model stagnates at a low compression rate and stops learning.</p>

<p>We tried other methods to resolve the degenerate distribution problem, including entropy penalties, but they were unstable, difficult to tune, and sometimes collapsed later in training. We ultimately solved the problem with brute force; any logit with absolute value greater than 4 gets a large l1 penalty applied to it. In practice, this means that the model is penalized for masking out a frame with probability under 2% or over 98%. We chose this threshold because in the worst case, around 13% of our compute would be wasted on duplicate latents, which was the most we were willing to tolerate. In practice, the amount of wasted compute was significantly lower (below 5%).<sup class="fn"><a id="fnref-5" href="#fn-5">5</a></sup></p>

<p>After training on a dataset of 200 million frames and 6k TPU v6e hours (which translates to around 3k H100 hours), we were able to get decent results with our 170M parameter model.</p>

<h3>Reconstruction Quality</h3>

<p>The VAE encodes 256x256 video into a compact latent (8x spatial compression: 768→96 channels per patch) and reconstructs it. Averaged over 600 real videos:</p>

<table>
  <thead>
    <tr><th>Metric</th><th class="num">Value</th></tr>
  </thead>
  <tbody>
    <tr><td>All-frames MSE</td><td class="num">0.0026</td></tr>
    <tr><td>Bernoulli-sampled MSE</td><td class="num">0.0025</td></tr>
    <tr><td>Mean frame keep rate</td><td class="num">18.5% (5.9 / 32 frames)</td></tr>
    <tr><td>Average temporal compression</td><td class="num">5.4x</td></tr>
    <tr><td>Combined compression (spatial × temporal)</td><td class="num">43x</td></tr>
  </tbody>
</table>

<figure>
  <img src="docs/recon_video0.gif" alt="Aerial coast — reconstruction">
  <figcaption>Aerial coast.</figcaption>
</figure>
<figure>
  <img src="docs/recon_video1.gif" alt="Beach jumping — reconstruction">
  <figcaption>Beach jumping.</figcaption>
</figure>
<figure>
  <img src="docs/recon_video2.gif" alt="Dance — high motion — reconstruction">
  <figcaption>Dance — high motion.</figcaption>
</figure>
<figure>
  <img src="docs/recon_video3.gif" alt="Music video — high motion — reconstruction">
  <figcaption>Music video — high motion.</figcaption>
</figure>

<h3>Frame-Budget Compression</h3>

<p>The encoder assigns each frame an importance score. We can keep only the top-K frames and fill the rest with a learned token, compressing at arbitrary temporal ratios on top of the 8x spatial compression. Each panel below shows, on a 32-frame clip: Original | Top-16 | Top-8 | Top-4 | Top-1.</p>

<figure>
  <img src="docs/compress_video0.gif" alt="Seascape — frame budget compression">
  <figcaption>Seascape.</figcaption>
</figure>
<figure>
  <img src="docs/compress_video1.gif" alt="Palm beach aerial — frame budget compression">
  <figcaption>Palm beach aerial.</figcaption>
</figure>
<figure>
  <img src="docs/compress_video2.gif" alt="Dance — high motion — frame budget compression">
  <figcaption>Dance — high motion.</figcaption>
</figure>
<figure>
  <img src="docs/compress_video3.gif" alt="Music video — high motion — frame budget compression">
  <figcaption>Music video — high motion.</figcaption>
</figure>

<p>To validate that the encoder was intelligently choosing frames, we inspected a few videos manually. As you can see, the slow moving videos are typically more compressed than the faster ones.</p>

<table>
  <thead>
    <tr>
      <th>Motion</th><th>Clip</th>
      <th class="num">Frames</th><th class="num">Kept</th>
      <th class="num">Keep rate</th><th class="num">Compression</th><th class="num">MSE</th>
    </tr>
  </thead>
  <tbody>
    <tr><td>slow</td><td>aerial coastline pan</td><td class="num">32</td><td class="num">9.2 ± 1.3</td><td class="num">28.7%</td><td class="num">3.5x</td><td class="num">0.0097</td></tr>
    <tr><td>slow</td><td>tropical island aerial</td><td class="num">32</td><td class="num">3.2 ± 1.3</td><td class="num">10.0%</td><td class="num">10.0x</td><td class="num">0.0011</td></tr>
    <tr><td>slow</td><td>seascape (static)</td><td class="num">32</td><td class="num">12.4 ± 1.3</td><td class="num">38.7%</td><td class="num">2.6x</td><td class="num">0.0024</td></tr>
    <tr><td>fast</td><td>beach jumping tricks</td><td class="num">32</td><td class="num">7.2 ± 1.0</td><td class="num">22.7%</td><td class="num">4.4x</td><td class="num">0.0028</td></tr>
    <tr><td>fast</td><td>dance performance</td><td class="num">32</td><td class="num">2.6 ± 0.9</td><td class="num">8.0%</td><td class="num">12.5x</td><td class="num">0.0006</td></tr>
    <tr><td>fast</td><td>music video</td><td class="num">32</td><td class="num">5.6 ± 0.9</td><td class="num">17.4%</td><td class="num">5.8x</td><td class="num">0.0016</td></tr>
    <tr><td>fast hands</td><td>cardistry (still scene)</td><td class="num">32</td><td class="num">6.6 ± 1.0</td><td class="num">20.5%</td><td class="num">4.9x</td><td class="num">0.0013</td></tr>
  </tbody>
</table>

<figure>
  <img src="docs/bernoulli_slow_aerial_coastline_pan_32f.gif" alt="Slow — aerial coastline pan">
  <figcaption>Slow — aerial coastline pan (32 frames).</figcaption>
</figure>
<figure>
  <img src="docs/bernoulli_slow_tropical_island_aerial_32f.gif" alt="Slow — tropical island aerial">
  <figcaption>Slow — tropical island aerial (32 frames).</figcaption>
</figure>
<figure>
  <img src="docs/bernoulli_slow_seascape_static_32f.gif" alt="Slow — seascape, static">
  <figcaption>Slow — seascape, static (32 frames).</figcaption>
</figure>
<figure>
  <img src="docs/bernoulli_fast_beach_jumping_tricks_32f.gif" alt="Fast — beach jumping tricks">
  <figcaption>Fast — beach jumping tricks (32 frames).</figcaption>
</figure>
<figure>
  <img src="docs/bernoulli_fast_dance_performance_32f.gif" alt="Fast — dance performance">
  <figcaption>Fast — dance performance (32 frames).</figcaption>
</figure>
<figure>
  <img src="docs/bernoulli_fast_music_video_32f.gif" alt="Fast — music video">
  <figcaption>Fast — music video (32 frames).</figcaption>
</figure>
<figure>
  <img src="docs/bernoulli_fasthands_cardistry_32f.gif" alt="Cardistry — still scene with fast hands">
  <figcaption>Cardistry — still scene with fast hand motion (32 frames).</figcaption>
</figure>

<h2>Diffusion Model</h2>

<p>At the end of the day, a VAE is only useful if it can create a good latent space for a generative model.<sup class="fn"><a id="fnref-6" href="#fn-6">6</a></sup></p>

<p>The diffusion model is not the focus of this blog post, so this blog won't go into the details of training it.</p>

<p>We didn't have enough remaining compute or data to properly train a proper diffusion<sup class="fn"><a id="fnref-7" href="#fn-7">7</a></sup> model, but we gave it our best shot anyways. After roughly 6k TPU v6e hours, here are the results for the 500M model:</p>

<figure>
  <img src="docs/generated_grid_4x4.gif" alt="4x4 generation grid">
  <figcaption>16 seeds, 8 latent frames each. Labels show seed and output frame count (28–33 frames).</figcaption>
</figure>

<p>This was my first time training a video model, so I'm not sure if these results are good for the training data or compute budget. It's almost certainly data starved and undertrained, these results seem decent enough for a general purpose video generation model under the given budget.</p>

<h2>Acknowledgements</h2>

<p>Thank you to Google's TRC program for providing us with the compute, and to Panda 70M for curating a large dataset.</p>

<div class="footnotes">
  <ol>
    <li id="fn-1">We omit the batch dimension for brevity.<a class="fn-back" href="#fnref-1">↩</a></li>
    <li id="fn-2">In this case, reconstruction loss refers to a combination of MSE, perceptual loss, and adversarial loss, as is standard for autoencoders; we collapse this into one term for clarity.<a class="fn-back" href="#fnref-2">↩</a></li>
    <li id="fn-3">If we had a half-frame, such as multiplying a frame by 0.5, the decoder would simply learn to amplify that signal, which removes any compression.<a class="fn-back" href="#fnref-3">↩</a></li>
    <li id="fn-4">The last subarray has a chance to be sampled equal to the sum of elements in that subarray. In this case, we have a 60% chance to sample an element from the last subarray.<a class="fn-back" href="#fnref-4">↩</a></li>
    <li id="fn-5">It's possible to resample if the second sample is identical to the first, but we trained on JAX with TPUs, and TPUs really don't like it when you do this.<a class="fn-back" href="#fnref-5">↩</a></li>
    <li id="fn-6">I might upset some people by saying this.<a class="fn-back" href="#fnref-6">↩</a></li>
    <li id="fn-7">Technically flow matching.<a class="fn-back" href="#fnref-7">↩</a></li>
  </ol>
</div>

</body>
</html>