How the Geometric Context Transformer Enables Streaming 3D Reconstruction in LingBot-Map

The Geometric Context Transformer (GCT) architecture achieves streaming 3D reconstruction by processing video frames through a causal transformer with paged KV-cache management, sliding-window attention, and 3D rotary positional embeddings, allowing constant GPU memory usage regardless of video length.

LingBot-Map implements the Geometric Context Transformer as a specialized neural architecture for online robotic mapping and 3D scene understanding. Unlike traditional batch-processing methods that require entire video sequences in memory, the GCT architecture processes frames sequentially while maintaining temporal geometric consistency through a carefully designed causal attention mechanism.

Core Architecture: The GCTBase Backbone

The GCTBase class in lingbot_map/models/gct_base.py defines the shared foundation for all GCT variants. This abstract base class establishes the prediction heads and forward-pass orchestration used by both batch and streaming implementations.

The base class initializes three primary prediction heads during construction:

class GCTBase(nn.Module, PyTorchModelHubMixin, ABC):
    def __init__(..., enable_camera=True, enable_depth=True, ...):
        self.camera_head = self._build_camera_head() if enable_camera else None
        self.depth_head  = self._build_depth_head()  if enable_depth else None
        self.point_head  = self._build_point_head()  if enable_point else None

These heads operate on shared token representations, enabling simultaneous prediction of camera poses, depth maps, and 3D world points from a single forward pass. The base class remains abstract until concrete subclasses provide an aggregator (the transformer encoder) and camera head implementation.

Streaming Implementation: GCTStream

The GCTStream class in lingbot_map/models/gct_stream.py extends GCTBase to enable online inference. It introduces streaming-specific hyperparameters including sliding-window size and 3D rotary positional embeddings (RoPE).

class GCTStream(GCTBase):
    def __init__(..., sliding_window_size=-1, enable_3d_rope=False, ...):
        self.sliding_window_size = sliding_window_size
        self.enable_3d_rope      = enable_3d_rope
        super().__init__(..., enable_3d_rope=enable_3d_rope, ...)

Building the Streaming Aggregator

The _build_aggregator method instantiates AggregatorStream, the causal transformer encoder that manages temporal state:

def _build_aggregator(self) -> nn.Module:
    return AggregatorStream(
        img_size=self.img_size,
        patch_size=self.patch_size,
        embed_dim=self.embed_dim,
        sliding_window_size=self.sliding_window_size,
        enable_3d_rope=self.enable_3d_rope,
        use_flashinfer=not self.use_sdpa,
        kv_cache_sliding_window=self.kv_cache_sliding_window,
        ...
    )

This aggregator creates separate frame-local and global transformer blocks, with the global blocks implementing cross-frame causal attention using a KV cache.

Building the Causal Camera Head

The streaming variant replaces the standard camera head with CameraCausalHead, which respects the same KV-cache policies as the aggregator:

def _build_camera_head(self) -> nn.Module:
    return CameraCausalHead(
        dim_in=2 * self.embed_dim,
        sliding_window_size=self.sliding_window_size,
        enable_3d_rope=self.enable_camera_3d_rope,
        rope_theta=self.camera_rope_theta,
        ...
    )

Causal Aggregation with KV-Cache Management

The AggregatorStream class in lingbot_map/aggregator/stream.py implements the core streaming logic. It maintains temporal consistency through two cache backends: FlashInfer for GPU-accelerated paged attention and SDPA (Scaled Dot Product Attention) as a pure-Python fallback.

Block Architecture

The aggregator constructs two types of transformer blocks:

1. Frame blocks: Handle intra-frame self-attention between spatial patches
2. Global blocks: Implement cross-frame causal attention using FlashInferBlock or SDPABlock

def _build_blocks(self, block_fn, depth, embed_dim, num_heads, ...):
    # Frame-local blocks

    self.frame_blocks = nn.ModuleList([
        block_fn(**block_params, rope=self.rope) for _ in range(depth)
    ])
    
    # Global blocks with KV cache support

    GlobalBlockCls = SDPABlock if self.use_sdpa else FlashInferBlock
    self.global_blocks = nn.ModuleList([
        GlobalBlockCls(**block_params,
                       kv_cache_sliding_window=self.kv_cache_sliding_window,
                       ...)
        for _ in range(depth)
    ])

KV-Cache Initialization

The aggregator initializes separate cache structures for each backend:

def _init_kv_cache(self):
    self.kv_cache_manager = None          # FlashInfer: created on first use

    self.kv_cache = {}                    # SDPA: dict-based tensor storage

    self.total_frames_processed = 0

The clean_kv_cache() method resets these structures when processing new videos, preventing memory leakage between sequences.

3D Rotary Positional Embeddings

When enable_3d_rope=True, the aggregator initializes WanRotaryPosEmbed to encode temporal frame indices into the attention mechanism:

def _init_3d_rope(self):
    self.rope3d = WanRotaryPosEmbed(
        attention_head_dim=head_dim,
        patch_size=(1, self.patch_size, self.patch_size),
        max_seq_len=self.max_frame_num,
    )

During streaming inference, the model generates position encodings using the global frame index (f_start, f_end), enabling the network to reason about temporal geometry across the video sequence.

Attention Blocks: FlashInfer vs. SDPA

The GCT architecture supports two attention implementations in lingbot_map/layers/block.py, allowing deployment across different hardware configurations.

FlashInferBlock (GPU-Accelerated)

FlashInferBlock uses paged KV cache management through FlashInferAttention. During streaming mode with num_frames <= 1, it pre-computes query/key/value tensors and manages cache appends:

if is_streaming:
    manager = kv_cache
    q_nhd, k_nhd, v_nhd = self.attn_pre(x, pos=pos, enable_3d_rope=enable_3d_rope)
    if skip_append:
        manager.append_frame(global_idx, k_nhd, v_nhd)
        attn_x = manager.compute_attention(global_idx, q_nhd)
        manager.rollback_last_frame(global_idx)
    else:
        manager.append_frame(global_idx, k_nhd, v_nhd)
        manager.evict_frames(...)
        attn_x = manager.compute_attention(global_idx, q_nhd)

The skip_append flag enables keyframe-only caching, where non-keyframes compute attention using existing cache entries without persisting their own KV tensors, drastically reducing memory growth.

SDPABlock (Fallback)

SDPABlock relies on PyTorch's native scaled_dot_product_attention with a dictionary-based KV cache. It implements identical logic to FlashInferBlock but without paged memory management, making it suitable for CPU inference or environments lacking FlashInfer support.

Pose Estimation and 3D Point Generation

The CameraCausalHead receives aggregated tokens and runs a small transformer that iteratively refines the 9-dimensional pose encoding. Like the aggregator, it respects KV-cache policies and can operate with or without storing KV tensors for non-keyframes.

From Encodings to World Points

After the forward pass, GCTBase provides utilities to convert pose encodings and depth predictions into world-space point clouds:

def _unproject_depth_to_world(self, depth, pose_enc):
    # Decode intrinsics / extrinsics from pose_enc

    extrinsics, intrinsics = pose_encoding_to_extri_intri(pose_enc, ...)
    
    # Build camera-to-world matrix

    c2w = closed_form_inverse_se3(extrinsics_4x4)
    
    # Back-project depth using intrinsics, then transform to world coordinates

    return world_points_h[..., :3]

This method, utilizing closed_form_inverse_se3 from the geometry utilities, enables real-time export of reconstructed 3D scenes during streaming.

Streaming Inference API

The inference_streaming method in GCTStream provides the high-level interface for processing long videos:

def inference_streaming(self, images, num_scale_frames=None,
                        keyframe_interval=1, output_device=None):
    # 1. Normalize input shape

    # 2. Process scale frames together (bidirectional attention)

    # 3. Stream remaining frames with causal attention

    # 4. Manage KV cache via _set_skip_append for non-keyframes

    # 5. Concatenate per-frame predictions

This method processes initial scale frames jointly using bidirectional attention to establish metric scale, then switches to causal streaming for subsequent frames.

Practical Usage Example

import torch
from lingbot_map.models.gct_stream import GCTStream

# Initialize streaming model with 3D RoPE and sliding window

model = GCTStream(
    img_size=518,
    patch_size=14,
    embed_dim=1024,
    enable_3d_rope=True,
    sliding_window_size=32,        # Limit attention to recent 32 frames

    enable_camera=True,
    enable_depth=True,
).cuda()

# Prepare streaming video input

video = torch.rand(100, 3, 518, 518).cuda()

# Run streaming inference with keyframe selection

output = model.inference_streaming(
    video,
    num_scale_frames=4,            # First 4 frames establish scale

    keyframe_interval=5,             # Cache KV every 5th frame only

    output_device=torch.device('cpu')
)

# Extract results

pose_enc = output["pose_enc"]          # Camera poses (1, 100, 9)

depth = output["depth"]                # Depth maps (1, 100, 518, 518, 1)

world_points = output["world_points"]  # 3D points (1, 100, 518, 518, 3)

This pipeline maintains constant GPU memory usage regardless of the 100-frame input length by evicting old KV cache entries and skipping storage for non-keyframes.

Summary

• GCTBase provides the shared backbone with prediction heads for camera, depth, and point outputs, defined in lingbot_map/models/gct_base.py.
• GCTStream extends the base with causal streaming capabilities, implementing inference_streaming with KV-cache management in lingbot_map/models/gct_stream.py.
• AggregatorStream maintains temporal consistency through FlashInferBlock or SDPABlock, supporting sliding-window attention and 3D rotary positional embeddings.
• Keyframe skipping via _skip_append and clean_kv_cache ensures memory consumption scales with the attention window rather than video length.
• CameraCausalHead and geometry utilities convert token representations into world-space point clouds suitable for real-time 3D reconstruction.

Frequently Asked Questions

How does the Geometric Context Transformer maintain constant memory usage during streaming?

The GCT architecture implements a sliding-window KV cache that only stores key and value tensors for recent frames or selected keyframes. By setting _skip_append=True for non-keyframes in FlashInferBlock or SDPABlock, the model computes attention using existing cache entries without persisting new tensors. The clean_kv_cache() method resets caches between videos, ensuring memory consumption remains bounded by the window size rather than the total video length.

What is the difference between FlashInfer and SDPA attention blocks in LingBot-Map?

FlashInferBlock utilizes paged GPU memory management through FlashInferAttention for high-throughput streaming inference, while SDPABlock uses PyTorch's native scaled_dot_product_attention with a dictionary-based cache. Both blocks implement identical causal masking and support 3D rotary positional embeddings, but FlashInfer offers superior performance for long sequences through optimized memory paging. The model automatically selects FlashInfer unless use_sdpa=True is specified.

How does 3D RoPE improve temporal consistency in streaming reconstruction?

3D Rotary Positional Embeddings (RoPE), implemented via WanRotaryPosEmbed in the aggregator, encode the absolute frame index into the attention mechanism using temporal coordinates. This allows the transformer to distinguish between frames based on their position in the video sequence, enabling explicit reasoning about temporal geometry and motion. When enable_3d_rope=True, the model uses the global frame range (f_start, f_end) to generate unique position encodings for each timestep.

Why are scale frames processed differently from streaming frames?

Scale frames require bidirectional attention to establish the metric scale of the scene before causal streaming begins. The inference_streaming method processes num_scale_frames jointly using the scale token, allowing the model to correlate distant features for accurate scale estimation. Subsequent frames switch to causal attention, where each frame only attends to itself and previous frames, enabling the memory-efficient streaming behavior necessary for long video sequences.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}