RWKV Seed Generator for Scene Completion

• RWKV-SG is a specialized module that transforms partial 3D point clouds into coarse, feature-rich outputs using a linear RWKV-based mechanism.
• It employs a modular architecture—with PointNet encoding, PRWKV stacks, and RWKV-ATTN—to effectively fuse local and global features while enhancing computational efficiency.
• Empirical evaluations demonstrate that RWKV-SG improves completion accuracy by about 25% in Chamfer Distance and significantly reduces model size compared to conventional baselines.

The RWKV Seed Generator (RWKV-SG) constitutes a specialized module for generating coarse, feature-rich point clouds from partial input data within the context of point cloud semantic scene completion. It is central to the architecture of RWKV-PCSSC, leveraging the Receptance Weighted Key Value (RWKV) mechanism to improve parameter and memory efficiency while delivering competitive or superior accuracy. RWKV-SG operates exclusively on geometry, eschewing auxiliary modalities such as color or normal vectors.

1. Architectural Overview

RWKV-SG transforms an input partial point cloud $P_\text{in} \in \mathbb{R}^{N\times3}$ into a coarse, completed point set with associated semantic logits and features. The module is highly modular, with each sub-block processing tensors of defined shape:

Stage	Input Shape/Type	Output Shape/Type
PointNet Encoding	$P_\text{in}\in\mathbb{R}^{N\times3}$	$F_0\in\mathbb{R}^{N\times C_f}$
PRWKV Stack (4 layers)	$F_{0}\in\mathbb{R}^{N\times C_f}$	$F_{in}\in\mathbb{R}^{N\times C_f}$
Global Feature (SA)	$F_{in}\in\mathbb{R}^{N\times C_f}$	$f\in\mathbb{R}^{C_g}$
Query Generation	$[P_{in}\Vert f]\in\mathbb{R}^{N\times(3+C_g)}$	$q_{in}\in\mathbb{R}^{N\times C_q}$
RWKV-ATTN	$P_{in}, q_{in}, k_{in}$	$P_\text{in}\in\mathbb{R}^{N\times3}$0
Deconvolution	$P_\text{in}\in\mathbb{R}^{N\times3}$1	$P_\text{in}\in\mathbb{R}^{N\times3}$2
Rebuild Head	$P_\text{in}\in\mathbb{R}^{N\times3}$3	$P_\text{in}\in\mathbb{R}^{N\times3}$4
Coarse Point Sampling	$P_\text{in}\in\mathbb{R}^{N\times3}$5	$P_\text{in}\in\mathbb{R}^{N\times3}$6
Segment Head	$P_\text{in}\in\mathbb{R}^{N\times3}$7	$P_\text{in}\in\mathbb{R}^{N\times3}$8

Following the preliminary feature extraction $P_\text{in}\in\mathbb{R}^{N\times3}$9 via a PointNet-style encoder, a four-layer PointRWKV (PRWKV) stack abstracts context from local and global neighborhoods, producing $F_0\in\mathbb{R}^{N\times C_f}$0. Global context $F_0\in\mathbb{R}^{N\times C_f}$1 is pooled from $F_0\in\mathbb{R}^{N\times C_f}$2 and broadcast to each point to form queries $F_0\in\mathbb{R}^{N\times C_f}$3 using an MLP. RWKV-ATTN fuses queries, keys, and spatial neighborhoods to estimate missing-region features $F_0\in\mathbb{R}^{N\times C_f}$4, which are deconvolved and reparameterized as position offsets $F_0\in\mathbb{R}^{N\times C_f}$5. Sampled coarse points $F_0\in\mathbb{R}^{N\times C_f}$6 and their features $F_0\in\mathbb{R}^{N\times C_f}$7 form the output, fed to a semantic segmentation head for coarse per-point class logits $F_0\in\mathbb{R}^{N\times C_f}$8.

2. Core RWKV Mechanism and Equations

RWKV modules substitute the quadratic $F_0\in\mathbb{R}^{N\times C_f}$9 softmax self-attention with linear-complexity “Receptance Weighted Key-Value” (RWKV) aggregation. For input point features $F_{0}\in\mathbb{R}^{N\times C_f}$0:

• P-Shift (per-channel local reordering):

$F_{0}\in\mathbb{R}^{N\times C_f}$1

($F_{0}\in\mathbb{R}^{N\times C_f}$2, $F_{0}\in\mathbb{R}^{N\times C_f}$3, $F_{0}\in\mathbb{R}^{N\times C_f}$4 are learnable.)

• Linear Projections:

$F_{0}\in\mathbb{R}^{N\times C_f}$5

• Bidirectional Linear Attention (Bi-WKV):

For output index $F_{0}\in\mathbb{R}^{N\times C_f}$6,

$F_{0}\in\mathbb{R}^{N\times C_f}$7

$F_{0}\in\mathbb{R}^{N\times C_f}$8

where $F_{0}\in\mathbb{R}^{N\times C_f}$9 are learnable scalars.

• Receptance Gating:

$F_{in}\in\mathbb{R}^{N\times C_f}$0

• Output:

$F_{in}\in\mathbb{R}^{N\times C_f}$1

Within RWKV-ATTN, a hybrid of global PRWKV output and local $F_{in}\in\mathbb{R}^{N\times C_f}$2-NN attention is used: - Local values: $F_{in}\in\mathbb{R}^{N\times C_f}$3 - Gated: $F_{in}\in\mathbb{R}^{N\times C_f}$4 - Weights: $F_{in}\in\mathbb{R}^{N\times C_f}$5 - Output: $F_{in}\in\mathbb{R}^{N\times C_f}$6

This structure enables global context aggregation with linear complexity and maintains spatial discrimination through local attention.

3. Feature Aggregation Workflow

The processing steps of RWKV-SG are as follows:

1. Preliminary Feature Extraction: $F_{in}\in\mathbb{R}^{N\times C_f}$7.
2. Contextual Abstraction: $F_{in}\in\mathbb{R}^{N\times C_f}$8 computed via four PRWKV layers.
3. Global Context Gathering: $F_{in}\in\mathbb{R}^{N\times C_f}$9 pooled using Set Abstraction; combined per-point with $F_{in}\in\mathbb{R}^{N\times C_f}$0 and processed to $F_{in}\in\mathbb{R}^{N\times C_f}$1.
4. Local and Global Feature Fusion: $F_{in}\in\mathbb{R}^{N\times C_f}$2 set to $F_{in}\in\mathbb{R}^{N\times C_f}$3; RWKV-ATTN computes $F_{in}\in\mathbb{R}^{N\times C_f}$4 per point within each $F_{in}\in\mathbb{R}^{N\times C_f}$5-NN neighborhood.
5. Missing Feature Deconvolution: $F_{in}\in\mathbb{R}^{N\times C_f}$6 upsampled to $F_{in}\in\mathbb{R}^{N\times C_f}$7 through a Snowflake-style deconvolution.
6. Coarse Completion: $F_{in}\in\mathbb{R}^{N\times C_f}$8 is regressed; $F_{in}\in\mathbb{R}^{N\times C_f}$9.
7. Farthest Point Sampling: $f\in\mathbb{R}^{C_g}$0 and $f\in\mathbb{R}^{C_g}$1 sampled to $f\in\mathbb{R}^{C_g}$2 coarse points.
8. Coarse Semantic Segmentation: Per-point logits $f\in\mathbb{R}^{C_g}$3 computed from $f\in\mathbb{R}^{C_g}$4.

This pipeline delivers plausible coarse geometry and features filling input holes, while maintaining efficiency through linear mechanisms.

4. Learnable Parameters and Model Efficiency

RWKV-SG is parameterized for compactness and speed. For typical feature dimensionality $f\in\mathbb{R}^{C_g}$5:

• PointNet encoder: $f\in\mathbb{R}^{C_g}$641K parameters.
• PRWKV stack (4 layers): $f\in\mathbb{R}^{C_g}$7M.
• Query-generation MLP: $f\in\mathbb{R}^{C_g}$80.26M.
• RWKV-ATTN internals: $f\in\mathbb{R}^{C_g}$90.40M.
• Deconvolution: $[P_{in}\Vert f]\in\mathbb{R}^{N\times(3+C_g)}$00.15M.
• Rebuild head: $[P_{in}\Vert f]\in\mathbb{R}^{N\times(3+C_g)}$10.03M.
• Segment head: $[P_{in}\Vert f]\in\mathbb{R}^{N\times(3+C_g)}$20.6M.

Total: $[P_{in}\Vert f]\in\mathbb{R}^{N\times(3+C_g)}$32.5M parameters, accounting for 50–60% of the full RWKV-PCSSC network. The linear-complexity RWKV structure enables the entire model (RWKV-SG + RWKV-PD) to remain $[P_{in}\Vert f]\in\mathbb{R}^{N\times(3+C_g)}$44M parameters—yielding a $[P_{in}\Vert f]\in\mathbb{R}^{N\times(3+C_g)}$5 reduction relative to the PointSSC baseline ($[P_{in}\Vert f]\in\mathbb{R}^{N\times(3+C_g)}$617M).

5. Empirical Performance and Impact

RWKV-SG and its accompanying network modules offer significant improvements in parameter and memory efficiency over softmax-attention-based dense architectures:

• Parameter reduction: Full model size $[P_{in}\Vert f]\in\mathbb{R}^{N\times(3+C_g)}$776.1% smaller than PointSSC.
• Memory efficiency: Peak GPU memory reduced by $[P_{in}\Vert f]\in\mathbb{R}^{N\times(3+C_g)}$827% (training, batch size 8, RTX3090).
• Ablation study: Removal of RWKV-SG in SSC-PC increases Chamfer Distance from 0.265 to 0.353 ($[P_{in}\Vert f]\in\mathbb{R}^{N\times(3+C_g)}$9 worse) and lowers mean accuracy from 97.99% to 97.49%.
• Qualitative output: $q_{in}\in\mathbb{R}^{N\times C_q}$0 clouds generated by RWKV-SG already reconstruct large missing areas plausibly.
• Downstream effect: RWKV-PD refinements act primarily on edges; RWKV-SG provides the structural estimate.
• Overall effect: RWKV-SG improves completion by $q_{in}\in\mathbb{R}^{N\times C_q}$125% in Chamfer Distance over non-RWKV baselines while preserving or exceeding state-of-the-art SSC accuracy.

A plausible implication is that the majority of completion accuracy and efficiency gains in RWKV-PCSSC can be attributed directly to the design of RWKV-SG.

6. Context and Significance within Point Cloud Completion

RWKV-SG exemplifies a new paradigm in 3D point cloud completion, replacing resource-intensive attention with a linear, context-aware mechanism. By forgoing auxiliary cues (color, normals) and reducing overparameterization, it delivers competitive semantic scene completion on both standard datasets (SSC-PC, NYUCAD-PC, PointSSC) and new benchmarks (NYUCAD-PC-V2, 3D-FRONT-PC), as developed in "RWKV-PCSSC: Exploring RWKV Model for Point Cloud Semantic Scene Completion" (He et al., 13 Nov 2025). This suggests further investigation of RWKV-style mechanisms is warranted for scalable 3D scene understanding in memory- and compute-constrained environments.